Update app.py

app.py

@@ -7,9 +7,12 @@ 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 replacement for LLMChain
+from langchain_core.runnables import RunnableSequence  # modern
 from langchain.tools import StructuredTool
 
+IMG_HEIGHT = 256
+IMG_WIDTH = 256
+
 # === Download model if not exists ===
 model_path = "unet_model.h5"
 if not os.path.exists(model_path):
@@ -24,36 +27,52 @@ if not os.path.exists(model_path):
 print("Loading model...")
 model = tf.keras.models.load_model(model_path, compile=False)
 
-# === Segmentation ===
-def classify_image(image_input):
-    img = tf.image.resize(image_input, (256, 256))
+
+# === Segmentation + Stats ===
+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)
 
-    prediction = model.predict(img)[0]  # (256, 256, 1) maybe
-
-    mask = (prediction > 0.5).astype(np.uint8) * 255  # binary mask
+    prediction = model.predict(img)[0]  # (256, 256, 1)
+    mask = (prediction > 0.5).astype(np.uint8) * 255
 
-    # Squeeze to (H, W) if needed
     if mask.ndim == 3 and mask.shape[-1] == 1:
         mask = np.squeeze(mask, axis=-1)
 
-    return mask
+    # Tumor stats
+    tumor_area = np.sum(prediction > 0.5)
+    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"
+
+    stats = {
+        "tumor_area": int(tumor_area),
+        "total_area": total_area,
+        "tumor_ratio": tumor_ratio,
+        "tumor_label": tumor_label
+    }
 
-# === Main handler ===
+    return mask, stats
+
+
+# === Gradio handler ===
 def rishigpt_handler(image_input, groq_api_key):
     os.environ["GROQ_API_KEY"] = groq_api_key
 
-    mask = classify_image(image_input)
+    mask, stats = classify_image_and_stats(image_input)
 
-    # Dummy tool for LangChain agent
-    def segment_brain_tool():
-        return "A brain tumor mask was generated."
+    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}."
 
     tool = StructuredTool.from_function(
         segment_brain_tool,
         name="segment_brain",
-        description="Segment brain MRI for tumor detection."
+        description="Provide tumor segmentation stats for the MRI image. Takes dummy input text."
     )
 
     llm = ChatGroq(
@@ -68,20 +87,23 @@ def rishigpt_handler(image_input, groq_api_key):
         verbose=True
     )
 
-    user_query = "I uploaded a brain MRI. What does the segmentation say?"
+    user_query = "Get segmentation details"
     classification = agent.run(user_query)
 
-    # New style: RunnableSequence
     prompt = PromptTemplate(
         input_variables=["result"],
-        template="You are a medical imaging expert. Based on the result: {result}, explain what this means for diagnosis."
+        template=(
+            "You are a medical imaging expert. Based on this tumor analysis result: {result}, "
+            "explain what this means for the patient in simple language."
+        )
     )
-    chain = prompt | llm
 
+    chain = prompt | llm
     description = chain.invoke({"result": classification})
 
     return mask, description
 
+
 # === Gradio UI ===
 inputs = [
     gr.Image(type="numpy", label="Upload Brain MRI Slice"),
@@ -99,5 +121,5 @@ if __name__ == "__main__":
         inputs=inputs,
         outputs=outputs,
         title="RishiGPT Medical Brain Segmentation",
-        description="UNet++ Brain Tumor Segmentation with LangChain integration"
+        description="UNet++ Brain Tumor Segmentation with LangChain integration. Includes tumor stats and medical explanation."
     ).launch()