Update app.py

app.py

@@ -1,3 +1,7 @@
+## Some code was modified from Ovseg and OV-Sam.Thanks to their excellent work.
+## Ovseg Code:https://github.com/facebookresearch/ov-seg
+## OV-Sam Code:https://github.com/HarborYuan/ovsam
+
 import spaces
 import multiprocessing as mp
 import numpy as np
@@ -18,6 +22,17 @@ import gradio as gr
 import open_clip
 from sam2.build_sam import build_sam2
 from mask_adapter.modeling.meta_arch.mask_adapter_head import build_mask_adapter
+from mask_adapter.data.datasets import openseg_classes
+
+COCO_CATEGORIES_pan = openseg_classes.get_coco_categories_with_prompt_eng()
+thing_classes = [k["name"] for k in COCO_CATEGORIES_pan if k["isthing"] == 1]
+stuff_classes = [k["name"] for k in COCO_CATEGORIES_pan]
+ADE20K_150_CATEGORIES_ = openseg_classes.get_ade20k_categories_with_prompt_eng()
+ade20k_thing_classes = [k["name"] for k in ADE20K_150_CATEGORIES_ if k["isthing"] == 1]
+ade20k_stuff_classes = [k["name"] for k in ADE20K_150_CATEGORIES_]
+class_names_coco_ade20k = thing_classes + stuff_classes + ade20k_thing_classes+ ade20k_stuff_classes
+
+
 
 
 
@@ -93,7 +108,7 @@ def inference_automatic(input_img, class_names):
 @spaces.GPU
 @torch.no_grad()
 @torch.autocast(device_type="cuda", dtype=torch.float32)
-def inference_point(input_img,  img_state,):
+def inference_point(input_img, img_state,class_names_input):
 
 
     mp.set_start_method("spawn", force=True)
@@ -106,8 +121,20 @@ def inference_point(input_img,  img_state,):
     
     demo = SAMPointVisualizationDemo(cfg, 0.8, sam2_model, clip_model,mask_adapter)
     
-    text_features = torch.from_numpy(np.load("./text_embedding/lvis_coco_text_embedding.npy")).cuda()
-    _, visualized_output = demo.run_on_image_with_points(img_state.img, points,text_features)
+    if not class_names_input:
+        class_names_input = class_names_coco_ade20k
+        
+    if class_names_input == class_names_coco_ade20k:
+        text_features = torch.from_numpy(np.load("./text_embedding/coco_ade20k_text_embedding.npy")).cuda()
+        _, visualized_output = demo.run_on_image_with_points(img_state.img, points,text_features)
+    else:
+        class_names_input = class_names_input.split(',')
+        txts = [f'a photo of {cls_name}' for cls_name in class_names_input]
+        text = open_clip.tokenize(txts)
+        text_features = clip_model.encode_text(text.cuda())
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+        _, visualized_output = demo.run_on_image_with_points(img_state.img, points,text_features,class_names_input)
+        
     return visualized_output
 
 
@@ -136,8 +163,20 @@ def inference_box(input_img,  img_state,):
     
     demo = SAMPointVisualizationDemo(cfg, 0.8, sam2_model, clip_model,mask_adapter)
     
-    text_features = torch.from_numpy(np.load("./text_embedding/lvis_coco_text_embedding.npy")).cuda()
-    _, visualized_output = demo.run_on_image_with_boxes(img_state.img, bbox,text_features)
+    if not class_names_input:
+        class_names_input = class_names_coco_ade20k
+        
+    if class_names_input == class_names_coco_ade20k:
+        text_features = torch.from_numpy(np.load("./text_embedding/coco_ade20k_text_embedding.npy")).cuda()
+        _, visualized_output = demo.run_on_image_with_boxes(img_state.img, bbox,text_features)
+    else:
+        class_names_input = class_names_input.split(',')
+        txts = [f'a photo of {cls_name}' for cls_name in class_names_input]
+        text = open_clip.tokenize(txts)
+        text_features = clip_model.encode_text(text.cuda())
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+        _, visualized_output = demo.run_on_image_with_boxes(img_state.img, bbox,text_features,class_names_input)
+        
     return visualized_output
 
 
@@ -234,7 +273,7 @@ def preprocess_example(input_img, img_state):
 
 def clear_everything(img_state):
     img_state.clear()
-    return img_state, None, None
+    return img_state, None, None, gr.Textbox(value='',lines=1, placeholder=class_names_coco_ade20k, label='Class Names')
 
 
 def clean_prompts(img_state):
@@ -296,7 +335,7 @@ with gr.Blocks() as demo:
                     output_image = gr.Image(type="pil", label='Segmentation Map')
 
                     # Buttons below segmentation map (now placed under segmentation map)
-                    run_button = gr.Button("Run Automatic Segmentation")
+                    run_button = gr.Button("Run Automatic Segmentation", elem_id="run_button",variant='primary')
                     run_button.click(inference_automatic, inputs=[input_image, class_names], outputs=output_image)
                     
                     clear_button = gr.Button("Clear")
@@ -310,9 +349,12 @@ with gr.Blocks() as demo:
             with gr.Row():  # 水平排列
                 with gr.Column(scale=1): 
                     input_image = gr.Image( label="Input Image", type="pil") 
-                with gr.Column(scale=1):  # 第二列：分割图输出
+                    class_names_input_box = gr.Textbox(lines=1, placeholder=class_names_coco_ade20k, label='Class Names')
+                with gr.Column(scale=1): 
                     output_image_box = gr.Image(type="pil", label='Segmentation Map',interactive=False)  # 输出分割图
-            
+                    clear_prompt_button_box = gr.Button("Clean Prompt")
+                    clear_button_box = gr.Button("Restart")
+                    
             gr.Markdown("Click the top-left and bottom-right corners of the image to select a rectangular area") 
 
             input_image.select(
@@ -321,30 +363,31 @@ with gr.Blocks() as demo:
                 outputs=[img_state_bbox, input_image]
             ).then(
                 inference_box,
-                inputs=[input_image, img_state_bbox],
+                inputs=[input_image, img_state_bbox,class_names_input_box],
                 outputs=[output_image_box]
             )
-            clear_prompt_button_box = gr.Button("Clean Prompt")
+            
+            
             clear_prompt_button_box.click(
                 clean_prompts,
                 inputs=[img_state_bbox],
                 outputs=[img_state_bbox, input_image, output_image_box]
             )
-            clear_button_box = gr.Button("Restart")
+            
             clear_button_box.click(
                 clear_everything,
                 inputs=[img_state_bbox],
-                outputs=[img_state_bbox, input_image, output_image_box]
+                outputs=[img_state_bbox, input_image, output_image_box,class_names_input_box]
             )
             input_image.clear(
                 clear_everything,
                 inputs=[img_state_bbox],
-                outputs=[img_state_bbox, input_image, output_image_box]
+                outputs=[img_state_bbox, input_image, output_image_box,class_names_input_box]
             )
             output_image_box.clear(
                 clear_everything,
                 inputs=[img_state_bbox],
-                outputs=[img_state_bbox, input_image, output_image_box]
+                outputs=[img_state_bbox, input_image, output_image_box,class_names_input_box]
             )
 
             
@@ -363,44 +406,41 @@ with gr.Blocks() as demo:
             with gr.Row():  # 水平排列
                 with gr.Column(scale=1): 
                     input_image = gr.Image( label="Input Image", type="pil") 
-                with gr.Column(scale=1):  # 第二列：分割图输出
+                    class_names_input_point = gr.Textbox(lines=1, placeholder=class_names_coco_ade20k, label='Class Names')
+                with gr.Column(scale=1):  
                     output_image_point = gr.Image(type="pil", label='Segmentation Map',interactive=False)  # 输出分割图
-
+                    clear_prompt_button_point = gr.Button("Clean Prompt")
+                    clear_button_point = gr.Button("Restart")
+                    
             input_image.select(
                 get_points_with_draw,
                 [input_image, img_state_points],
                 outputs=[img_state_points, input_image]
             ).then(
                 inference_point,
-                inputs=[input_image, img_state_points],
+                inputs=[input_image, img_state_points,class_names_input_point],
                 outputs=[output_image_point]
             )
-            clear_prompt_button_point = gr.Button("Clean Prompt")
             clear_prompt_button_point.click(
                 clean_prompts,
                 inputs=[img_state_points],
                 outputs=[img_state_points, input_image, output_image_point]
             )
-            clear_button_point = gr.Button("Restart")
             clear_button_point.click(
                 clear_everything,
                 inputs=[img_state_points],
-                outputs=[img_state_points, input_image, output_image_point]
+                outputs=[img_state_points, input_image, output_image_point,class_names_input_point]
             )
             input_image.clear(
                 clear_everything,
                 inputs=[img_state_points],
-                outputs=[img_state_points, input_image, output_image_point]
+                outputs=[img_state_points, input_image, output_image_point,class_names_input_point]
             )
             output_image_point.clear(
                 clear_everything,
                 inputs=[img_state_points],
-                outputs=[img_state_points, input_image, output_image_point]
+                outputs=[img_state_points, input_image, output_image_point,class_names_input_point]
             )
-            def clear_and_set_example_point(example):
-                clear_everything(img_state_points)
-                return example
-
             gr.Examples(
                     examples=examples_point, 
                     inputs=[input_image, img_state_points],