Update app.py

app.py

@@ -23,6 +23,7 @@ torch.manual_seed(0)
 
 CFG_PATH = "configs/demo/pokemon.yaml"
 
+
 def generate_distinct_colors(n):
     colors = []
     # generate a random number from 0 to 1
@@ -79,7 +80,6 @@ def visualize_segmentation(image,
     # Create a figure and axis
     fig, ax = plt.subplots(1, figsize=(12, 9))
     # Display the image
-    # ax.imshow(image)
     # Generate distinct colors for each mask
     final_mask = np.zeros(
         (masks.shape[1], masks.shape[2], 3), dtype=np.float32)
@@ -115,18 +115,11 @@ def get_sam_masks(cfg,
                   image_path=None,
                   img_sam=None,
                   pipeline=None):
-    # image_id = image_path.split('/')[-1].split('.')[0]
-    # sam_mask_path = os.path.join(cfg.test.sam_mask_root, f'{image_id}.npz')
-    # if os.path.exists(sam_mask_path):
-    #     sam_mask_masks = np.load(sam_mask_path, allow_pickle=True)
-    #     mask_tensor = torch.from_numpy(sam_mask_masks['mask_tensor'])
-    #     mask_list = sam_mask_path['mask_list']
-    # else:
     print("generating sam masks online")
     if img_sam is None and image_path is not None:
         raise ValueError(
             'Please provide either the image path or the image numpy array.')
-    
+
     mask_tensor, mask_list = generate_masks_from_sam(
         image_path,
         save_path='./',
@@ -168,6 +161,7 @@ def load_sam(cfg, device):
     )
     return pipeline
 
+
 def generate(img,
              class_names,
              clip_thresh,
@@ -192,7 +186,6 @@ def generate(img,
                     seg_mode='semantic',
                     device=device)
 
- 
     # resize image by dividing 2 if the size is larger than 1000
     if img.size[0] > 1000:
         img = img.resize((img.size[0] // 2, img.size[1] // 2))
@@ -203,8 +196,7 @@ def generate(img,
 
     # class_names = ['the women chatting', 'the women chatting', 'table', 'fridge', 'cooking pot']
 
-    pseudo_masks, _, _ = car_model(
-        img, sentences, 1)
+    pseudo_masks, _ = car_model(img, sentences)
 
     if post_process == 'SAM':
         pipeline = load_sam(cfg, device)
@@ -230,7 +222,6 @@ def generate(img,
     return demo_img
 
 
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser('car')
     parser.add_argument("--cfg-path",
@@ -238,48 +229,48 @@ if __name__ == "__main__":
                         help="path to configuration file.")
     args = parser.parse_args()
 
-    demo = gr.Interface(generate, 
-        inputs=[gr.Image(label="upload an image", type="pil"), 
-                "text",
-                gr.Slider(label="clip thresh",
-                          minimum=0,
-                          maximum=1,
-                          value=0.4,
-                          step=0.1,
-                info="the threshold for clip-es adversarial heatmap clipping"),
-                gr.Slider(label="mask thresh",
-                          minimum=0,
-                          maximum=1,
-                          value=0.6,
-                          step=0.1,
-                info="the binariation threshold for the mask to generate visual prompt"),
-                gr.Slider(label="confidence thresh",
-                          minimum=0,
-                          maximum=1,
-                          value=0,
-                          step=0.1,
-                info="the threshold for filtering the proposed classes"),
-                gr.Radio(["CRF", "SAM"], label="post process", value="CRF", info="choose the post process method"),
-                gr.Slider(label="stability score thresh for SAM mask proposal \n(only when SAM is chosen for post process)",
-                          minimum=0,
-                          maximum=1,
-                          value=0.95,
-                          step=0.1),
-                gr.Slider(label="box nms thresh for SAM mask proposal \n(only when SAM is chosen for post process)", minimum=0, maximum=1, value=0.7, step=0.1),
-                gr.Slider(label="intersection over mask threshold for SAM mask proposal \n(only when SAM is chosen for post process)", minimum=0, maximum=1, value=0.5, step=0.1),
-                gr.Slider(label="minimum prediction threshold for SAM mask proposal \n(only when SAM is chosen for post process)", minimum=0, maximum=1, value=0.03, step=0.01)],
-        outputs="image",
-        title="CLIP as RNN: Segment Countless Visual Concepts without Training Endeavor",
-        description="This is the official demo for CLIP as RNN. Please upload an image and type in the class names (connected by ',' e.g. cat,dog,human) you want to segment. The model will generate the segmentation masks for the input image. You can also adjust the clip thresh, mask thresh and confidence thresh to get better results.",
-        examples=[["demo/pokemon1.jpg", "Charmander,Bulbasaur,Squirtle", 0.6, 0.6, 0, "SAM", 0.95, 0.7, 0.6, 0.01],
-                    ["demo/batman.jpg", "Batman,Joker,Cat Woman", 0.6, 0.6, 0, "SAM", 0.95, 0.7, 0.6, 0.01],
-                    ["demo/avengers1.jpg", "Thor,Captain America,Hulk,Iron Man", 0.6, 0.6, 0, "SAM", 0.89, 0.65, 0.5, 0.03],
-
-                    ])
+    demo = gr.Interface(generate,
+                        inputs=[gr.Image(label="upload an image", type="pil"),
+                                "text",
+                                gr.Slider(label="clip thresh",
+                                          minimum=0,
+                                          maximum=1,
+                                          value=0.4,
+                                          step=0.1,
+                                          info="the threshold for clip-es adversarial heatmap clipping"),
+                                gr.Slider(label="mask thresh",
+                                          minimum=0,
+                                          maximum=1,
+                                          value=0.6,
+                                          step=0.1,
+                                          info="the binariation threshold for the mask to generate visual prompt"),
+                                gr.Slider(label="confidence thresh",
+                                          minimum=0,
+                                          maximum=1,
+                                          value=0,
+                                          step=0.1,
+                                          info="the threshold for filtering the proposed classes"),
+                                gr.Radio(["CRF", "SAM"], label="post process",
+                                         value="CRF", info="choose the post process method"),
+                                gr.Slider(label="stability score thresh for SAM mask proposal \n(only when SAM is chosen for post process)",
+                                          minimum=0,
+                                          maximum=1,
+                                          value=0.95,
+                                          step=0.1),
+                                gr.Slider(label="box nms thresh for SAM mask proposal \n(only when SAM is chosen for post process)",
+                                          minimum=0, maximum=1, value=0.7, step=0.1),
+                                gr.Slider(label="intersection over mask threshold for SAM mask proposal \n(only when SAM is chosen for post process)",
+                                          minimum=0, maximum=1, value=0.5, step=0.1),
+                                gr.Slider(label="minimum prediction threshold for SAM mask proposal \n(only when SAM is chosen for post process)", minimum=0, maximum=1, value=0.03, step=0.01)],
+                        outputs="image",
+                        title="CLIP as RNN: Segment Countless Visual Concepts without Training Endeavor",
+                        description="This is the official demo for CLIP as RNN. Please upload an image and type in the class names (connected by ',' e.g. cat,dog,human) you want to segment. The model will generate the segmentation masks for the input image. You can also adjust the clip thresh, mask thresh and confidence thresh to get better results.",
+                        examples=[["demo/pokemon.jpg", "Pikachu,Eevee", 0.6, 0.6, 0, "SAM", 0.95, 0.7, 0.6, 0.01],
+                                  ["demo/Eiffel_tower.jpg", "Eiffel Tower",
+                                   0.6, 0.6, 0, "SAM", 0.95, 0.7, 0.6, 0.01],
+                                  ["demo/superhero.jpeg", "Batman,Superman,Wonder Woman,Flash,Cyborg",
+                                   0.6, 0.6, 0, "SAM", 0.89, 0.65, 0.5, 0.03],
+                                  ])
     demo.launch(share=True)
 
-    
     # device = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-stop = 0