Restore GPU version - remove CPU optimizations and restore GPU-compatible dependencies

app.py

@@ -3,6 +3,11 @@ from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, Generat
 from qwen_vl_utils import process_vision_info
 import torch
 import requests
+from ultralytics import YOLO
+from PIL import Image
+import matplotlib.pyplot as plt
+import numpy as np
+import io
 
 # ----------------------------
 # MODEL LOADING (MedVLM-R1) - CPU Compatible
@@ -31,6 +36,30 @@ temp_generation_config = GenerationConfig(
     pad_token_id=151643,
 )
 
+# ----------------------------
+# YOLO MODEL LOADING
+# ----------------------------
+yolo_model = YOLO("MedSegVLM_with_DeepSeek/best.pt")  # replace with your segmentation weights
+
+def inference(image_path: str):
+    """Runs YOLO segmentation on an image and returns the annotated image."""
+    # Load image
+    img = Image.open(image_path).convert("RGB")
+
+    # Run inference
+    results = yolo_model(img)
+
+    # Plot with masks and bounding boxes
+    annotated = results[0].plot()  # NumPy array (BGR)
+
+    # Convert from BGR (OpenCV default) to RGB for matplotlib
+    annotated_rgb = annotated[:, :, ::-1]
+
+    # Convert numpy array to PIL Image
+    annotated_image = Image.fromarray(annotated_rgb)
+    
+    return annotated_image
+
 # ----------------------------
 # API SETTINGS (DeepSeek R1)
 # ----------------------------
@@ -54,7 +83,10 @@ Your task:
 # ----------------------------
 def process_pipeline(image, user_question):
     if image is None or user_question.strip() == "":
-        return "Please upload an image and enter a question."
+        return "Please upload an image and enter a question.", None
+
+    # Run YOLO inference and get segmented image
+    segmented_image = inference(image)
 
     # Combine user's question with default
     combined_question = user_question.strip() + "\n\n" + DEFAULT_QUESTION
@@ -135,9 +167,9 @@ Original Answer:
     try:
         detailed_answer = response.json()["choices"][0]["message"]["content"]
     except Exception as e:
-        return f"**Error from DeepSeek:** {str(e)}\n\n```\n{response.text}\n```"
+        return f"**Error from DeepSeek:** {str(e)}\n\n```\n{response.text}\n```", segmented_image
 
-    return f"{detailed_answer}"
+    return f"{detailed_answer}", segmented_image
 
 
 # ----------------------------
@@ -145,31 +177,33 @@ Original Answer:
 # ----------------------------
 with gr.Blocks(title="Brain MRI QA") as demo:
     with gr.Row():
-        # Left column
+        # First column: input image and result image side by side
         with gr.Column():
-            image_input = gr.Image(type="filepath", label="Upload Medical Image")
+            with gr.Row():
+                image_input = gr.Image(type="filepath", label="Upload Medical Image")
+                result_image = gr.Image(type="filepath", label="Upload Medical Image")  # next to input image
+
             question_box = gr.Textbox(
                 label="Your Question about the Image",
                 placeholder="Type your question here..."
             )
-            # Buttons side by side
             with gr.Row():
                 submit_btn = gr.Button("Submit")
                 clear_btn = gr.Button("Clear")
 
-        # Right column
+        # Second column: LLM answer output
         with gr.Column():
             llm_output = gr.Markdown(label="Detailed LLM Answer")
-
     submit_btn.click(
         fn=process_pipeline,
         inputs=[image_input, question_box],
-        outputs=llm_output
+        outputs=[llm_output, result_image]
     )
     clear_btn.click(
-        fn=lambda: ("", ""),
-        outputs=[question_box, llm_output]
+        fn=lambda: ("", "", None),
+        outputs=[question_box, llm_output, result_image]
     )
 
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:148e1be81643a7f594d2dd626a7a27bab6b4fd03f197dde4e1b2005a578e37a3
+size 5468627

requirements.txt

@@ -1,3 +1,17 @@
+# gradio==5.42.0
+# transformers>=4.40.0
+# torch>=2.0.0
+# torchvision>=0.15.0
+# requests>=2.31.0
+# Pillow>=10.0.0
+# accelerate>=0.20.0
+# safetensors>=0.3.0
+# tokenizers>=0.15.0
+# numpy>=1.24.0
+# scipy>=1.10.0
+# qwen-vl-utils
+# ipython>=8.0.0
+
 gradio==5.42.0
 transformers>=4.40.0
 torch>=2.0.0
@@ -11,3 +25,6 @@ numpy>=1.24.0
 scipy>=1.10.0
 qwen-vl-utils
 ipython>=8.0.0
+ultralytics>=8.0.0
+matplotlib>=3.5.0
+opencv-python>=4.5.0