Update app.py

app.py

@@ -79,11 +79,11 @@ import torch
 # -----------------------------
 # 1. Load models
 # -----------------------------
-yolo_model = YOLO("best.pt")  # put your YOLO model in root directory
+yolo_model = YOLO("best.pt")  # make sure 'best.pt' is uploaded to root
 chatbot = pipeline("text-generation", model="microsoft/BioGPT-Large")
 
 # -----------------------------
-# 2. Define the logic
+# 2. Define main logic
 # -----------------------------
 def medical_chat(image, question):
     if image is None:
@@ -94,38 +94,40 @@ def medical_chat(image, question):
     predicted_class = results[0].names[int(results[0].probs.top1)]
     yolo_output = f"The YOLO model classified the image as: **{predicted_class}**."
 
-    # Direct YOLO query
+    # Direct YOLO question
     if "output of yolo" in question.lower():
         return yolo_output
 
-    # Explanation via LLM
+    # Explanation using LLM
     prompt = f"The YOLO model detected {predicted_class}. {question}"
     response = chatbot(prompt, max_new_tokens=200)
     answer = response[0]['generated_text']
     return f"{yolo_output}\n\n**Explanation:** {answer}"
 
 # -----------------------------
-# 3. Build a Gradio Interface (not ChatInterface)
+# 3. Build custom UI using Blocks (forces correct layout)
 # -----------------------------
-demo = gr.Interface(
-    fn=medical_chat,
-    inputs=[
-        gr.Image(type="pil", label="🩻 Upload Medical Image"),
-        gr.Textbox(label="💬 Ask your question", placeholder="e.g., What is the output of YOLO? or Explain this disease.")
-    ],
-    outputs="text",
-    title="🧠 YOLO + Medical Chatbot",
-    description="Upload a medical image → YOLO detects the disease → Chatbot explains or answers questions.",
-    examples=[
-        ["sample_xray.jpg", "What is the output of YOLO?"],
-        ["sample_xray.jpg", "Explain this disease."]
-    ]
-)
+with gr.Blocks(title="🧠 YOLO + Medical Chatbot") as demo:
+    gr.Markdown("## 🩺 Medical Image Analyzer & Chatbot")
+    gr.Markdown("Upload an image → YOLO detects disease → Chatbot explains it.")
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="🩻 Upload Medical Image")
+            text_input = gr.Textbox(
+                label="💬 Ask your question",
+                placeholder="e.g., What is the output of YOLO? or Explain this disease."
+            )
+            submit_btn = gr.Button("Analyze & Ask")
+
+        with gr.Column():
+            output_text = gr.Textbox(label="🧠 Chatbot Response", lines=10)
+
+    # When user clicks button
+    submit_btn.click(fn=medical_chat, inputs=[image_input, text_input], outputs=output_text)
 
 # -----------------------------
-# 4. Launch
+# 4. Launch the app
 # -----------------------------
 if __name__ == "__main__":
-    demo.launch()
-
-
+    demo.launch()