Update app.py

app.py

@@ -7,9 +7,10 @@ import requests
 from langchain_groq import ChatGroq
 from langchain.agents import initialize_agent
 from langchain.prompts import PromptTemplate
-from langchain_core.runnables import RunnableSequence  # modern
+from langchain_core.runnables import RunnableSequence
 from langchain.tools import StructuredTool
 
+
 IMG_HEIGHT = 256
 IMG_WIDTH = 256
 
@@ -28,27 +29,31 @@ print("Loading model...")
 model = tf.keras.models.load_model(model_path, compile=False)
 
 
-# === Segmentation + Stats ===
+# === Segmentation + Stats + Overlay ===
 def classify_image_and_stats(image_input):
     img = tf.image.resize(image_input, [IMG_HEIGHT, IMG_WIDTH])
-    img = img / 255.0
-    img = np.expand_dims(img, axis=0)
+    img_norm = img / 255.0
+    img_batch = np.expand_dims(img_norm, axis=0)
 
-    prediction = model.predict(img)[0]  # (256, 256, 1)
-    mask = (prediction > 0.5).astype(np.uint8) * 255
+    prediction = model.predict(img_batch)[0]  # (256, 256, 1)
+    mask = (prediction > 0.5).astype(np.uint8)
 
     if mask.ndim == 3 and mask.shape[-1] == 1:
         mask = np.squeeze(mask, axis=-1)
 
     # Tumor stats
-    tumor_area = np.sum(prediction > 0.5)
+    tumor_area = np.sum(mask)
     total_area = IMG_HEIGHT * IMG_WIDTH
     tumor_ratio = tumor_area / total_area
 
-    if tumor_ratio > 0.01:
-        tumor_label = "Tumor Detected"
-    else:
-        tumor_label = "No Tumor Detected"
+    tumor_label = "Tumor Detected" if tumor_ratio > 0.005 else "No Tumor Detected"
+
+    # === Overlay mask on original ===
+    overlay = np.array(img)  # original resized input
+    red_mask = np.zeros_like(overlay)
+    red_mask[..., 0] = mask * 255  # Red channel
+
+    overlay_img = np.clip(0.6 * overlay + 0.4 * red_mask, 0, 255).astype(np.uint8)
 
     stats = {
         "tumor_area": int(tumor_area),
@@ -57,51 +62,58 @@ def classify_image_and_stats(image_input):
         "tumor_label": tumor_label
     }
 
-    return mask, stats
+    return overlay_img, stats
 
 
 # === Gradio handler ===
 def rishigpt_handler(image_input, groq_api_key):
     os.environ["GROQ_API_KEY"] = groq_api_key
 
-    mask, stats = classify_image_and_stats(image_input)
+    overlay_img, stats = classify_image_and_stats(image_input)
 
     def segment_brain_tool(input_text: str) -> str:
-        return f"Tumor label: {stats['tumor_label']}. Tumor area: {stats['tumor_area']}. Ratio: {stats['tumor_ratio']:.4f}."
+        return (
+            f"Tumor label: {stats['tumor_label']}. "
+            f"Tumor area: {stats['tumor_area']}. "
+            f"Ratio: {stats['tumor_ratio']:.4f}."
+        )
 
     tool = StructuredTool.from_function(
         segment_brain_tool,
         name="segment_brain",
-        description="Provide tumor segmentation stats for the MRI image. Takes dummy input text."
+        description="Provide tumor segmentation stats for the MRI image."
     )
 
     llm = ChatGroq(
         model="meta-llama/llama-4-scout-17b-16e-instruct",
-        temperature=0.3
+        temperature=0.4
     )
 
     agent = initialize_agent(
         tools=[tool],
         llm=llm,
         agent="zero-shot-react-description",
-        verbose=True
+        verbose=False
     )
 
-    user_query = "Get segmentation details"
+    user_query = "Give me the segmentation details"
     classification = agent.run(user_query)
 
+    # Better prompt + output parser
     prompt = PromptTemplate(
         input_variables=["result"],
         template=(
-            "You are a medical imaging expert. Based on this tumor analysis result: {result}, "
-            "explain what this means for the patient in simple language."
+            "You are a compassionate AI radiologist. "
+            "Read this tumor analysis result: {result}. "
+            "Summarize the situation like you're talking to the patient in calm, clear language. "
+            "Add any recommendations for next steps too, but keep it easy to understand."
         )
     )
 
     chain = prompt | llm
-    description = chain.invoke({"result": classification})
+    description = chain.invoke({"result": classification}).content.strip()
 
-    return mask, description
+    return overlay_img, description
 
 
 # === Gradio UI ===
@@ -111,8 +123,8 @@ inputs = [
 ]
 
 outputs = [
-    gr.Image(type="numpy", label="Tumor Segmentation Mask"),
-    gr.Textbox(label="Medical Explanation")
+    gr.Image(type="numpy", label="Overlay: Brain MRI + Tumor Mask"),
+    gr.Textbox(label="Doctor's Explanation")
 ]
 
 if __name__ == "__main__":
@@ -120,6 +132,6 @@ if __name__ == "__main__":
         fn=rishigpt_handler,
         inputs=inputs,
         outputs=outputs,
-        title="RishiGPT Medical Brain Segmentation",
-        description="UNet++ Brain Tumor Segmentation with LangChain integration. Includes tumor stats and medical explanation."
+        title="🧠 RishiGPT Medical Brain Segmentation",
+        description="UNet++ Brain Tumor Segmentation with mask overlay, detailed stats, and human-like explanation."
     ).launch()