Update app.py

app.py

@@ -3,8 +3,16 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from PIL import Image 
 import numpy as np 
 import os 
+import tempfile
 import gradio as gr
 
+import cv2
+try:
+    from mmengine.visualization import Visualizer
+except ImportError:
+    Visualizer = None
+    print("Warning: mmengine is not installed, visualization is disabled.")
+    
 # Load the model and tokenizer 
 model_path = "ByteDance/Sa2VA-4B"
  
@@ -20,6 +28,17 @@ tokenizer = AutoTokenizer.from_pretrained(
     trust_remote_code = True,
 )
 
+def visualize(pred_mask, image_path, work_dir):
+    visualizer = Visualizer()
+    img = cv2.imread(image_path)
+    visualizer.set_image(img)
+    visualizer.draw_binary_masks(pred_mask, colors='g', alphas=0.4)
+    visual_result = visualizer.get_image()
+
+    output_path = os.path.join(work_dir, os.path.basename(image_path))
+    cv2.imwrite(output_path, visual_result)
+    return output_path
+
 def image_vision(image_input_path, prompt):
     image_path = image_input_path
     text_prompts = f"<image>{prompt}"
@@ -34,6 +53,7 @@ def image_vision(image_input_path, prompt):
     return_dict = model.predict_forward(**input_dict)
     print(return_dict)
     answer = return_dict["prediction"] # the text format answer
+    
     seg_image = return_dict["prediction_masks"]
     
     return answer, seg_image
@@ -41,7 +61,15 @@ def image_vision(image_input_path, prompt):
 def main_infer(image_input_path, prompt):
 
     answer, seg_image = image_vision(image_input_path, prompt)
-    return answer, seg_image[0]
+    pred_masks = seg_image[0]
+    
+    if '[SEG]' in answer and Visualizer is not None:
+        temp_dir = tempfile.mkdtemp()
+        pred_mask = pred_masks[0]
+        os.makedirs(temp_dir, exist_ok=True)
+        seg_result = visualize(pred_mask, image_input_path, temp_dir)
+    
+    return answer, seg_result
 
 # Gradio UI
 
@@ -56,7 +84,7 @@ with gr.Blocks() as demo:
                     submit_btn = gr.Button("Submit", scale=1)
             with gr.Column():
                 output_res = gr.Textbox(label="Response")
-                output_image = gr.Image(label="Segmentation")
+                output_image = gr.Image(label="Segmentation", type="numpy")
 
     submit_btn.click(
         fn = main_infer,