update app.py

app.py

@@ -434,8 +434,7 @@ def concatenate_images_vertical(image1, image2):
 
     return new_image
 
-def relate_anything(input_image_mask, k):    
-    input_image = input_image_mask['image']
+def relate_anything(input_image, k):    
     logger.info(f'relate_anything_1_{input_image.size}_')
     w, h = input_image.size
     max_edge = 1500
@@ -478,15 +477,17 @@ def relate_anything(input_image_mask, k):
         concate_pil_image = concatenate_images_vertical(current_pil_image, title_image)
         pil_image_list.append(concate_pil_image)
 
-    logger.info(f'relate_anything_5_')
-    yield pil_image_list
-
+    return pil_image_list
 
 mask_source_draw = "draw a mask on input image"
 mask_source_segment = "type what to detect below"
 
 def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, 
             iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation):
+    if (task_type == 'relate anything'):
+        output_images = relate_anything(input_image['image'], num_relation)
+        return output_images, gr.Gallery.update(label='relate images')
+
     text_prompt = text_prompt.strip()
     if not ((task_type == 'inpainting' or task_type == 'remove') and mask_source_radio == mask_source_draw):
         if text_prompt == '':
@@ -510,7 +511,7 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
     size = image_pil.size
 
     output_images = []
-    # output_images.append(input_image['image'])
+    output_images.append(input_image['image'])
     # run grounding dino model
     if (task_type == 'inpainting' or task_type == 'remove') and mask_source_radio == mask_source_draw:
         pass
@@ -538,11 +539,12 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
             "labels": pred_phrases,
         }
         image_with_box = plot_boxes_to_image(copy.deepcopy(image_pil), pred_dict)[0]
-        image_path = os.path.join(output_dir, f"grounding_dino_output_{file_temp}.jpg")
-        image_with_box.save(image_path)
-        detection_image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
-        os.remove(image_path)
-        output_images.append(detection_image_result)
+        # image_path = os.path.join(output_dir, f"grounding_dino_output_{file_temp}.jpg")
+        # image_with_box.save(image_path)
+        # detection_image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+        # os.remove(image_path)
+        # output_images.append(detection_image_result)
+        output_images.append(image_with_box)
 
     logger.info(f'run_anything_task_[{file_temp}]_{task_type}_2_')
     if task_type == 'segment' or ((task_type == 'inpainting' or task_type == 'remove') and mask_source_radio == mask_source_segment):
@@ -600,13 +602,12 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
             mask = masks[0][0].cpu().numpy()
             mask_pil = Image.fromarray(mask)   
 
-        image_path = os.path.join(output_dir, f"image_mask_{file_temp}.jpg")
-        # if reverse_mask:
-        #     mask_pil = mask_pil.point(lambda _: 255-_)
-        mask_pil.convert("RGB").save(image_path)
-        image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
-        os.remove(image_path)
-        output_images.append(image_result)
+        # image_path = os.path.join(output_dir, f"image_mask_{file_temp}.jpg")
+        # mask_pil.convert("RGB").save(image_path)
+        # image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+        # os.remove(image_path)
+        # output_images.append(image_result)
+        output_images.append(mask_pil.convert("RGB"))
 
         if task_type == 'inpainting':
             # inpainting pipeline
@@ -645,24 +646,23 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
                         mask_imgs.append(mask_pil_exp)
                 mask_pil = mix_masks(mask_imgs)
 
-                image_path = os.path.join(output_dir, f"image_mask_{file_temp}.jpg")
-                # if reverse_mask:
-                #     mask_pil = mask_pil.point(lambda _: 255-_)
-                mask_pil.convert("RGB").save(image_path)
-                image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
-                os.remove(image_path)
-                output_images.append(image_result)               
+                # image_path = os.path.join(output_dir, f"image_mask_{file_temp}.jpg")
+                # mask_pil.convert("RGB").save(image_path)
+                # image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+                # os.remove(image_path)
+                # output_images.append(image_result)
+                output_images.append(mask_pil.convert("RGB"))               
             image_inpainting = lama_cleaner_process(np.array(image_pil), np.array(mask_pil.convert("L")))
 
         image_inpainting = image_inpainting.resize((image_pil.size[0], image_pil.size[1]))
 
-        image_path = os.path.join(output_dir, f"grounded_sam_inpainting_output_{file_temp}.jpg")
-        image_inpainting.save(image_path)
-        image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
-        os.remove(image_path)
-        logger.info(f'run_anything_task_[{file_temp}]_{task_type}_9__{type(image_inpainting)} / {type(image_result)}')
-        output_images.append(image_inpainting)
+        # image_path = os.path.join(output_dir, f"grounded_sam_inpainting_output_{file_temp}.jpg")
+        # image_inpainting.save(image_path)
+        # image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+        # os.remove(image_path)
+        # logger.info(f'run_anything_task_[{file_temp}]_{task_type}_9_')
         # output_images.append(image_result)
+        output_images.append(image_inpainting)
         return output_images, gr.Gallery.update(label='result images')        
     else:
         logger.info(f"task_type:{task_type} error!")
@@ -674,10 +674,10 @@ def change_radio_display(task_type, mask_source_radio):
     inpaint_prompt_visible = False
     mask_source_radio_visible = False
     num_relation_visible = False
-    run_button_visible = True
-    relate_all_button_visible = False
-    gsa_gallery_visible = True 
-    ram_gallery_visible = False
+    # run_button_visible = True
+    # relate_all_button_visible = False
+    # gsa_gallery_visible = True 
+    # ram_gallery_visible = False
     if task_type == "inpainting":
         inpaint_prompt_visible = True
     if task_type == "inpainting" or task_type == "remove":
@@ -687,11 +687,12 @@ def change_radio_display(task_type, mask_source_radio):
     if task_type == "relate anything":
         text_prompt_visible = False
         num_relation_visible = True
-        run_button_visible = False
-        relate_all_button_visible = True
-        gsa_gallery_visible = False 
-        ram_gallery_visible = True             
-    return  gr.Textbox.update(visible=text_prompt_visible), gr.Textbox.update(visible=inpaint_prompt_visible), gr.Radio.update(visible=mask_source_radio_visible), gr.Slider.update(visible=num_relation_visible), gr.Button.update(visible=run_button_visible), gr.Button.update(visible=relate_all_button_visible), gr.Gallery.update(visible=gsa_gallery_visible), gr.Gallery.update(visible=ram_gallery_visible)
+        # run_button_visible = False
+        # relate_all_button_visible = True
+        # gsa_gallery_visible = False 
+        # ram_gallery_visible = True             
+    return  gr.Textbox.update(visible=text_prompt_visible), gr.Textbox.update(visible=inpaint_prompt_visible), gr.Radio.update(visible=mask_source_radio_visible), gr.Slider.update(visible=num_relation_visible) 
+    #, gr.Button.update(visible=run_button_visible), gr.Button.update(visible=relate_all_button_visible), gr.Gallery.update(visible=gsa_gallery_visible), gr.Gallery.update(visible=ram_gallery_visible)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser("Grounded SAM demo", add_help=True)
@@ -715,7 +716,7 @@ if __name__ == "__main__":
                 inpaint_prompt = gr.Textbox(label="Inpaint Prompt (if this is empty, then remove)", visible=False)
                 num_relation = gr.Slider(label="How many relations do you want to see", minimum=1, maximum=20, value=5, step=1, visible=False)
                 run_button = gr.Button(label="Run", visible=True)
-                relate_all_button = gr.Button(label="Run", visible=False)
+                # relate_all_button = gr.Button(label="Run", visible=False)
                 with gr.Accordion("Advanced options", open=False) as advanced_options:
                     box_threshold = gr.Slider(
                         label="Box Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.001
@@ -734,17 +735,19 @@ if __name__ == "__main__":
                             remove_mask_extend = gr.Textbox(label="remove_mask_extend", value='10')
 
             with gr.Column():
-                gsa_gallery = gr.Gallery(label="result images", show_label=True, elem_id="gsa_allery", visible=True
-                    ).style(preview=True, grid=[2], full_width=True, full_height=True)
-                ram_gallery = gr.Gallery(label="Your Result", show_label=True, elem_id="ram_gallery", visible=False
-                    ).style(preview=True, columns=5, object_fit="scale-down")
+                image_gallery = gr.Gallery(label="result images", show_label=True, elem_id="gsa_allery", visible=True
+                    ).style(preview=True, columns=[5], object_fit="scale-down", height="auto")                
+                # gsa_gallery = gr.Gallery(label="result images", show_label=True, elem_id="gsa_allery", visible=True
+                #     ).style(preview=True, grid=[2], full_width=True, full_height=True)
+                # ram_gallery = gr.Gallery(label="Your Result", show_label=True, elem_id="ram_gallery", visible=False
+                #     ).style(preview=True, columns=5, object_fit="scale-down")
 
         run_button.click(fn=run_anything_task, inputs=[
-                        input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation], outputs=[gsa_gallery, gsa_gallery], show_progress=True, queue=True)
-        relate_all_button.click(fn=relate_anything, inputs=[input_image, num_relation], outputs=[ram_gallery], show_progress=True, queue=True)
+                        input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation], outputs=[image_gallery, image_gallery], show_progress=True, queue=True)
+        # relate_all_button.click(fn=relate_anything, inputs=[input_image, num_relation], outputs=[ram_gallery], show_progress=True, queue=True)
 
-        task_type.change(fn=change_radio_display, inputs=[task_type, mask_source_radio], outputs=[text_prompt, inpaint_prompt, mask_source_radio, num_relation, run_button, relate_all_button, gsa_gallery, ram_gallery])
-        mask_source_radio.change(fn=change_radio_display, inputs=[task_type, mask_source_radio], outputs=[text_prompt, inpaint_prompt, mask_source_radio, num_relation, run_button, relate_all_button, gsa_gallery, ram_gallery])
+        task_type.change(fn=change_radio_display, inputs=[task_type, mask_source_radio], outputs=[text_prompt, inpaint_prompt, mask_source_radio, num_relation])
+        mask_source_radio.change(fn=change_radio_display, inputs=[task_type, mask_source_radio], outputs=[text_prompt, inpaint_prompt, mask_source_radio, num_relation])
 
         DESCRIPTION = '### This demo from [Grounded-Segment-Anything](https://github.com/IDEA-Research/Grounded-Segment-Anything). <br>'
         DESCRIPTION += 'RAM from [RelateAnything](https://github.com/Luodian/RelateAnything). <br>'