Upload app.py

app.py

@@ -4,7 +4,7 @@ import warnings
 import copy
 os.system("python -m pip install -e asam")
 os.system("python -m pip install -e GroundingDINO")
-os.system("python -m pip uninstall gradio")
+# os.system("python -m pip uninstall gradio")
 os.system("python -m pip install gradio==3.38.0")
 os.system("pip install opencv-python pycocotools matplotlib onnxruntime onnx ipykernel")
 sys.path.append(os.path.join(os.getcwd(), "GroundingDINO"))
@@ -162,6 +162,7 @@ sam_predictor = None
 
 
 def run_grounded_sam(input_image, text_prompt, task_type, box_threshold, text_threshold, iou_threshold):    
+    print(text_prompt, type(text_prompt))
     global blip_processor, blip_model, groundingdino_model, sam_predictor
 
     # make dir
@@ -170,7 +171,8 @@ def run_grounded_sam(input_image, text_prompt, task_type, box_threshold, text_th
     scribble = np.array(input_image["mask"])
     image_pil = input_image["image"].convert("RGB")
     transformed_image = transform_image(image_pil)
-
+    print('img sum:' ,torch.sum(transformed_image).to(torch.int).item())
+    
     if groundingdino_model is None:
         groundingdino_model = load_model(
             config_file, ckpt_filenmae, device=device)
@@ -225,8 +227,34 @@ def run_grounded_sam(input_image, text_prompt, task_type, box_threshold, text_th
         # use NMS to handle overlapped boxes
         print(f"Revise caption with number: {text_prompt}")
     
-    if task_type == 'text' or task_type == 'automatic' or task_type == 'scribble_box':
-        if task_type == 'scribble_box':
+    if task_type == 'default_box' or task_type == 'automatic' or task_type == 'scribble_box':
+        if task_type == 'default_box':
+            id = torch.sum(transformed_image).to(torch.int).item()
+            if id == -1683627: #example 1 *
+                x_min, y_min, x_max, y_max = 204, 213, 813, 1023
+            elif id == 1137390: #example 2 *
+                x_min, y_min, x_max, y_max = 125, 168, 842, 904
+            elif id == 1145309: #example 3 *
+                x_min, y_min, x_max, y_max = 0, 486, 992, 899
+            elif id == 1091779: #example 4 *
+                x_min, y_min, x_max, y_max = 2, 73, 981, 968
+            elif id == -1335352: #example 5 *
+                x_min, y_min, x_max, y_max = 201, 195, 811, 1023
+            elif id == -1479645: #example 6
+                x_min, y_min, x_max, y_max = 428, 0, 992, 799
+            elif id == -544197: #example 7
+                x_min, y_min, x_max, y_max = 106, 419, 312, 783
+            elif id == -23873: #example 8
+                x_min, y_min, x_max, y_max = 250, 25, 774, 803
+            elif id == -1572157: #example 9 *
+                x_min, y_min, x_max, y_max = 15, 88, 1006, 977
+            else:
+                print("not defined")
+                raise NotImplementedError
+            bbox = np.array([x_min, y_min, x_max, y_max])
+            bbox = torch.tensor(bbox).unsqueeze(0)
+            transformed_boxes = sam_predictor.transform.apply_boxes_torch(bbox, image.shape[:2]).to(device)
+        elif task_type == 'scribble_box':
             scribble = scribble.transpose(2, 1, 0)[0]
             labeled_array, num_features = ndimage.label(scribble >= 255)
             centers = ndimage.center_of_mass(scribble, labeled_array, range(1, num_features+1))
@@ -261,7 +289,7 @@ def run_grounded_sam(input_image, text_prompt, task_type, box_threshold, text_th
             draw_mask(mask[0].cpu().numpy(), mask_draw, random_color=True)
         image_draw = ImageDraw.Draw(image_pil)
 
-        if task_type == 'scribble_box':
+        if task_type == 'scribble_box' or task_type == 'default_box':
             for box in bbox:
                 draw_box(box, image_draw, None)
         else:
@@ -305,7 +333,7 @@ def run_grounded_sam(input_image, text_prompt, task_type, box_threshold, text_th
             draw_mask(a_mask[0].cpu().numpy(), a_mask_draw, random_color=True)
         a_image_draw = ImageDraw.Draw(a_image_pil)
 
-        if task_type == 'scribble_box':
+        if task_type == 'scribble_box' or task_type == 'default_box':
             for box in bbox:
                 draw_box(box, a_image_draw, None)
         else:
@@ -373,8 +401,6 @@ def run_grounded_sam(input_image, text_prompt, task_type, box_threshold, text_th
         a_image_pil.alpha_composite(a_mask_image)
         
         return [[image_pil, mask_image],[a_image_pil, a_mask_image]]
-        
-    
 
     else:
         print("task_type:{} error!".format(task_type))
@@ -406,14 +432,14 @@ if __name__ == "__main__":
         You may check the instruction below, or check our github page about more details.
         <details>
         You may select an example image or upload your image to start, we support 4 prompt types:
-
+        
+        **default_box**: According to the mask label, automaticly generate the default box prompt, only used for examples.
+        
         **automatic**: Automaticly generate text prompt and the corresponding box input with BLIP and Grounding-DINO.
 
         **scribble_point**: Click an point on the target instance.
 
         **scribble_box**: Click on two points, the top-left point and the bottom-right point to represent a bounding box of the target instance.
-
-        **text**: Send text prompt to identify the target instance in the `Text prompt` box.
         
         </details>
         """)
@@ -421,14 +447,14 @@ if __name__ == "__main__":
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(
-                    source='upload', type="pil", value="example1.jpg", tool="sketch",brush_radius=20)
+                    source='upload', type="pil", value="example9.jpg", tool="sketch",brush_radius=20)
                 task_type = gr.Dropdown(
-                    ["automatic", "scribble_point", "scribble_box", "text"], value="automatic", label="task_type")
-                text_prompt = gr.Textbox(label="Text Prompt", placeholder="bench .")
+                    ["default_box","automatic", "scribble_point", "scribble_box", "text"], value="default_box", label="task_type")
+                text_prompt = gr.Textbox(label="Text Prompt", placeholder="bench .", visible=False)
                 run_button = gr.Button(label="Run")
                 with gr.Accordion("Advanced options", open=False):
                     box_threshold = gr.Slider(
-                        label="Box Threshold", minimum=0.0, maximum=1.0, value=0.2, step=0.001
+                        label="Box Threshold", minimum=0.0, maximum=1.0, value=0.4, step=0.001
                     )
                     text_threshold = gr.Slider(
                         label="Text Threshold", minimum=0.0, maximum=1.0, value=0.25, step=0.001
@@ -461,8 +487,19 @@ if __name__ == "__main__":
             with gr.Column():            
                 gr.Examples(["example2.jpg"], inputs=input_image)
             with gr.Column():
-                gr.Examples(["example3.jpg"], inputs=input_image)
-
+                 gr.Examples(["example3.jpg"], inputs=input_image)
+            with gr.Column():
+                gr.Examples(["example4.jpg"], inputs=input_image)
+            with gr.Column():
+                gr.Examples(["example5.jpg"], inputs=input_image)
+            with gr.Column():            
+                gr.Examples(["example6.jpg"], inputs=input_image)
+            with gr.Column():
+                gr.Examples(["example7.jpg"], inputs=input_image)
+            with gr.Column():
+                gr.Examples(["example8.jpg"], inputs=input_image)
+            with gr.Column():
+                gr.Examples(["example9.jpg"], inputs=input_image)            
         run_button.click(fn=run_grounded_sam, inputs=[
             input_image, text_prompt, task_type, box_threshold, text_threshold, iou_threshold], outputs=[gallery1,gallery2])