🚀 Added all models + Gradio app.py + requirements.txt

app.py

@@ -5,10 +5,10 @@ from torchvision import transforms
 from PIL import Image
 import google.generativeai as genai
 
-# Inject API key
+# ------------------ Gemini API Key (Hardcoded - as requested) ------------------
 genai.configure(api_key="AIzaSyDn5yK2_2pYMId3bpFlAf0LkWoJ7dvEcqM")
 
-# ------------------ Model Definitions ------------------
+# ------------------ CNN Model Definitions ------------------
 
 class AlzheimerCNN(torch.nn.Module):
     def __init__(self, num_classes=2):
@@ -18,6 +18,7 @@ class AlzheimerCNN(torch.nn.Module):
         self.pool = torch.nn.MaxPool2d(2, 2)
         self.fc1 = torch.nn.Linear(32 * 32 * 32, 64)
         self.fc2 = torch.nn.Linear(64, num_classes)
+
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
         x = self.pool(F.relu(self.conv2(x)))
@@ -34,6 +35,7 @@ class BrainHemorrhageCNN(torch.nn.Module):
         self.dropout = torch.nn.Dropout(0.5)
         self.fc1 = torch.nn.Linear(32 * 16 * 16, 64)
         self.fc2 = torch.nn.Linear(64, num_classes)
+
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
         x = self.pool(F.relu(self.conv2(x)))
@@ -50,6 +52,7 @@ class StrokeCTCNN(torch.nn.Module):
         self.pool = torch.nn.MaxPool2d(2, 2)
         self.fc1 = torch.nn.Linear(32 * 32 * 32, 64)
         self.fc2 = torch.nn.Linear(64, num_classes)
+
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
         x = self.pool(F.relu(self.conv2(x)))
@@ -65,6 +68,7 @@ class BrainTumorCNN(torch.nn.Module):
         self.pool = torch.nn.MaxPool2d(2, 2)
         self.fc1 = torch.nn.Linear(32 * 32 * 32, 64)
         self.fc2 = torch.nn.Linear(64, num_classes)
+
     def forward(self, x):
         x = self.pool(F.relu(self.conv1(x)))
         x = self.pool(F.relu(self.conv2(x)))
@@ -72,7 +76,7 @@ class BrainTumorCNN(torch.nn.Module):
         x = F.relu(self.fc1(x))
         return self.fc2(x)
 
-# ------------------ Load Models ------------------
+# ------------------ Model Loading ------------------
 
 transform = transforms.Compose([
     transforms.Resize((128, 128)),
@@ -83,34 +87,34 @@ transform = transforms.Compose([
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 alz_model = AlzheimerCNN()
-alz_model.load_state_dict(torch.load(r"C:\Users\HP\OneDrive\Desktop\trained model\alzmodel.pth", map_location=device))
+alz_model.load_state_dict(torch.load("alzmodel.pth", map_location=device))
 alz_model.eval()
 
 brainh_model = BrainHemorrhageCNN()
-brainh_model.load_state_dict(torch.load(r"C:\Users\HP\OneDrive\Desktop\trained model\brainham.pth", map_location=device))
+brainh_model.load_state_dict(torch.load("brainham.pth", map_location=device))
 brainh_model.eval()
 
 brainst_model = StrokeCTCNN()
-brainst_model.load_state_dict(torch.load(r"C:\Users\HP\OneDrive\Desktop\trained model\brainst_model.pth", map_location=device))
+brainst_model.load_state_dict(torch.load("brainst_model.pth", map_location=device))
 brainst_model.eval()
 
 braint_model = BrainTumorCNN()
-braint_model.load_state_dict(torch.load(r"C:\Users\HP\OneDrive\Desktop\trained model\braintmodel.pth", map_location=device))
+braint_model.load_state_dict(torch.load("braintmodel.pth", map_location=device))
 braint_model.eval()
 
-# ------------------ Recommendation Mapping ------------------
+# ------------------ Recommendations ------------------
 
 recommendation_dict = {
-    "Non Demented": "No signs of Alzheimer's detected. Maintain a healthy brain with mental activities and regular checkups(recomend you to use the gemini chatbot).",
-    "Demented": "Signs of dementia detected. Consult a neurologist for proper diagnosis and treatment options.(recomend you to use the gemini chatbot)",
-    "Normal": "Brain scan appears normal. Stay consistent with health checks.(recomend you to use the gemini chatbot)",
-    "Hemorrhagic": "Hemorrhage detected. Seek immediate medical attention — may require surgery or ICU.(recomend you to use the gemini chatbot)",
-    "Bleeding": "Bleeding stroke identified. Emergency treatment may be necessary.(recomend you to use the gemini chatbot)",
-    "Ischemia": "Ischemic stroke detected. Treatment may include clot-busting medication.(recomend you to use the gemini chatbot)",
-    "Glioma": "Glioma tumor found. Requires MRI follow-up and oncology consultation.(recomend you to use the gemini chatbot)",
-    "Meningioma": "Meningioma detected. Often benign but may need surgical evaluation.(recomend you to use the gemini chatbot)",
-    "No Tumor": "No brain tumor detected. Continue routine monitoring.(recomend you to use the gemini chatbot)",
-    "Pituitary": "Pituitary tumor detected. Hormonal and visual exams recommended.(recomend you to use the gemini chatbot)"
+    "Non Demented": "No signs of Alzheimer's detected. Maintain a healthy brain with mental activities and regular checkups.",
+    "Demented": "Signs of dementia detected. Consult a neurologist for proper diagnosis and treatment options.",
+    "Normal": "Brain scan appears normal. Stay consistent with health checks.",
+    "Hemorrhagic": "Hemorrhage detected. Seek immediate medical attention — may require surgery or ICU.",
+    "Bleeding": "Bleeding stroke identified. Emergency treatment may be necessary.",
+    "Ischemia": "Ischemic stroke detected. Treatment may include clot-busting medication.",
+    "Glioma": "Glioma tumor found. Requires MRI follow-up and oncology consultation.",
+    "Meningioma": "Meningioma detected. Often benign but may need surgical evaluation.",
+    "No Tumor": "No brain tumor detected. Continue routine monitoring.",
+    "Pituitary": "Pituitary tumor detected. Hormonal and visual exams recommended."
 }
 
 # ------------------ Prediction Function ------------------
@@ -122,49 +126,50 @@ def predict(disorder, image):
     if disorder == "Alzheimer":
         outputs = alz_model(img_tensor)
         class_names = ["Non Demented", "Demented"]
-
     elif disorder == "Brain Hemorrhage":
         outputs = brainh_model(img_tensor)
         class_names = ["Normal", "Hemorrhagic"]
-
     elif disorder == "Brain Stroke":
         outputs = brainst_model(img_tensor)
         class_names = ["Bleeding", "Ischemia", "Normal"]
-
     elif disorder == "Brain Tumor":
         outputs = braint_model(img_tensor)
         class_names = ["Glioma", "Meningioma", "No Tumor", "Pituitary"]
 
     probs = torch.softmax(outputs, dim=1)[0]
-    pred_idx = torch.argmax(probs).item()
-    pred_label = class_names[pred_idx]
+    pred_label = class_names[torch.argmax(probs).item()]
     recommendation = recommendation_dict.get(pred_label, "No recommendation available.")
 
-    return f"Prediction: {pred_label}\n\nRecommendation: {recommendation}"
+    return f"🧠 Prediction: {pred_label}\n\n📌 Recommendation: {recommendation}"
 
-# ------------------ Gemini Chatbot Function ------------------
+# ------------------ Gemini Chatbot ------------------
 
 model = genai.GenerativeModel("models/gemini-1.5-flash")
 
 def chat_with_gemini(user_input, history=[]):
-    prompt = "You are a highly experienced neurologist. Help answer questions related to brain diseases like stroke, dementia, tumors, Alzheimer’s, etc. and you answer questions short and simple to make a non medical field based person understand it better and easily\n"
-    convo = model.start_chat(history=[])
+    prompt = (
+        "You are a highly experienced and friendly neurologist. "
+        "Help answer questions related to brain diseases like stroke, dementia, tumors, Alzheimer’s, etc. "
+        "Answer briefly and simply so non-medical users can understand easily.\n\n"
+    )
+    convo = model.start_chat()
     convo.send_message(prompt + user_input)
     return convo.last.text
 
-# ------------------ Gradio Interface ------------------
+# ------------------ Gradio UI ------------------
 
 with gr.Blocks(theme=gr.themes.Base(), title="Brain Neurologist App") as demo:
-    with gr.Tab("Brain Scan Predictor"):
-        gr.Markdown("## 🧠 MRI/CT Brain Disorder Detection")
+    with gr.Tab("🧠 Brain Scan Predictor"):
+        gr.Markdown("## Upload an MRI/CT scan to get brain disorder prediction and medical advice")
         disorder_input = gr.Dropdown(["Alzheimer", "Brain Hemorrhage", "Brain Stroke", "Brain Tumor"], label="Select Disorder")
         image_input = gr.Image(type="filepath", label="Upload Brain Scan")
         output_text = gr.Textbox(label="Prediction and Recommendation")
         submit_btn = gr.Button("Predict")
         submit_btn.click(fn=predict, inputs=[disorder_input, image_input], outputs=output_text)
 
-    with gr.Tab("🧑‍⚕️ Dr.Neuro"):
-        gr.Markdown("## Gemini powered Neurologist Chatbot")
+    with gr.Tab("🧑‍⚕️ Dr.Neuro - Gemini Chatbot"):
+        gr.Markdown("## Ask neurological questions powered by Gemini")
         chatbot = gr.ChatInterface(fn=chat_with_gemini)
 
-demo.launch(share=True)
+if __name__ == "__main__":
+    demo.launch()