|
|
import gradio as gr |
|
|
import os |
|
|
import tensorflow as tf |
|
|
import numpy as np |
|
|
import requests |
|
|
|
|
|
from langchain_groq import ChatGroq |
|
|
from langchain.prompts import PromptTemplate |
|
|
from langchain.chains import LLMChain |
|
|
|
|
|
|
|
|
model_path = "unet_model.h5" |
|
|
|
|
|
|
|
|
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}...") |
|
|
with requests.get(hf_url, stream=True) as r: |
|
|
r.raise_for_status() |
|
|
with open(model_path, "wb") as f: |
|
|
for chunk in r.iter_content(chunk_size=8192): |
|
|
f.write(chunk) |
|
|
|
|
|
print("Loading model...") |
|
|
model = tf.keras.models.load_model(model_path, compile=False) |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
return mask |
|
|
|
|
|
|
|
|
def rishigpt_handler(image_input, groq_api_key): |
|
|
os.environ["GROQ_API_KEY"] = groq_api_key |
|
|
|
|
|
mask = classify_image(image_input) |
|
|
|
|
|
llm = ChatGroq( |
|
|
model="meta-llama/llama-4-scout-17b-16e-instruct", |
|
|
temperature=0.3 |
|
|
) |
|
|
|
|
|
prompt = PromptTemplate( |
|
|
input_variables=["result"], |
|
|
template="You are a medical imaging expert. Based on the result: {result}, explain what this means for diagnosis." |
|
|
) |
|
|
|
|
|
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}) |
|
|
|
|
|
return mask, description |
|
|
|
|
|
|
|
|
inputs = [ |
|
|
gr.Image(type="numpy", label="Upload Brain MRI Slice"), |
|
|
gr.Textbox(type="password", label="Groq API Key") |
|
|
] |
|
|
|
|
|
outputs = [ |
|
|
gr.Image(type="numpy", label="Tumor Segmentation Mask"), |
|
|
gr.Textbox(label="Medical Explanation") |
|
|
] |
|
|
|
|
|
if __name__ == "__main__": |
|
|
gr.Interface( |
|
|
fn=rishigpt_handler, |
|
|
inputs=inputs, |
|
|
outputs=outputs, |
|
|
title="RishiGPT Medical Brain Segmentation", |
|
|
description="UNet++ Brain Tumor Segmentation with LangChain integration" |
|
|
).launch() |
|
|
|
