Update app.py

app.py

@@ -5,13 +5,13 @@ import numpy as np
 import requests
 
 from langchain_groq import ChatGroq
+from langchain.agents import initialize_agent
 from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
-
+from langchain_core.runnables import RunnableSequence  # Modern replacement for LLMChain
+from langchain.tools import StructuredTool
 
+# === Download model if not exists ===
 model_path = "unet_model.h5"
-
-# Safe download with streaming to avoid incomplete file
 if not os.path.exists(model_path):
     hf_url = "https://huggingface.co/rishirajbal/UNET_plus_plus_Brain_segmentation/resolve/main/unet_model.h5"
     print(f"Downloading model from {hf_url}...")
@@ -24,44 +24,65 @@ 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))
     img = img / 255.0
     img = np.expand_dims(img, axis=0)
 
-    prediction = model.predict(img)[0]
-    mask = (prediction > 0.5).astype(np.uint8) * 255
+    prediction = model.predict(img)[0]  # (256, 256, 1) maybe
 
-    return mask
+    mask = (prediction > 0.5).astype(np.uint8) * 255  # binary mask
 
+    # Squeeze to (H, W) if needed
+    if mask.ndim == 3 and mask.shape[-1] == 1:
+        mask = np.squeeze(mask, axis=-1)
 
+    return mask
+
+# === Main handler ===
 def rishigpt_handler(image_input, groq_api_key):
     os.environ["GROQ_API_KEY"] = groq_api_key
 
     mask = classify_image(image_input)
 
+    # Dummy tool for LangChain agent
+    def segment_brain_tool():
+        return "A brain tumor mask was generated."
+
+    tool = StructuredTool.from_function(
+        segment_brain_tool,
+        name="segment_brain",
+        description="Segment brain MRI for tumor detection."
+    )
+
     llm = ChatGroq(
         model="meta-llama/llama-4-scout-17b-16e-instruct",
         temperature=0.3
     )
 
+    agent = initialize_agent(
+        tools=[tool],
+        llm=llm,
+        agent="zero-shot-react-description",
+        verbose=True
+    )
+
+    user_query = "I uploaded a brain MRI. What does the segmentation say?"
+    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."
     )
+    chain = prompt | llm
 
-    llm_chain = LLMChain(
-        llm=llm,
-        prompt=prompt
-    )
-
-    classification = "The brain tumor mask has been generated and segmentation is complete."
-    description = llm_chain.run({"result": classification})
+    description = chain.invoke({"result": classification})
 
     return mask, description
 
-
+# === Gradio UI ===
 inputs = [
     gr.Image(type="numpy", label="Upload Brain MRI Slice"),
     gr.Textbox(type="password", label="Groq API Key")