Update app.py

app.py

@@ -4,37 +4,47 @@ from PIL import Image
 from torchvision import transforms
 from TumorModel import TumorClassification, GliomaStageModel
 
-# ✅ Load tumor classification model
+# ─── Load Tumor Classification Model ───────────────────────────────────────────
 tumor_model = TumorClassification()
-tumor_model.load_state_dict(torch.load("BTD_model.pth", map_location=torch.device("cpu")))
+sd = torch.load("BTD_model.pth", map_location="cpu")
+renamed_sd = {}
+for k, v in sd.items():
+    new_key = (k
+        .replace("con1d.", "model.0.")
+        .replace("con2d.", "model.3.")
+        .replace("con3d.", "model.6.")
+        .replace("fc1.", "model.8.")
+        .replace("fc2.", "model.10.")
+        .replace("output.", "model.12."))
+    renamed_sd[new_key] = v
+tumor_model.load_state_dict(renamed_sd)
 tumor_model.eval()
 
-# ✅ Load glioma stage classification model
+# ─── Load Glioma Stage Model ───────────────────────────────────────────────────
 glioma_model = GliomaStageModel()
-glioma_model.load_state_dict(torch.load("glioma_stages.pth", map_location=torch.device("cpu")))
+glioma_model.load_state_dict(torch.load("glioma_stages.pth", map_location="cpu"))
 glioma_model.eval()
 
-# ✅ Labels
+# ─── Labels and Image Transform ───────────────────────────────────────────────
 tumor_labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
-stage_labels = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
+stage_labels = ['Stage 1', 'Stage 2']  # Or adjust to match your second model
 
-# ✅ Transform (resize to 208x208 to match training)
 transform = transforms.Compose([
     transforms.Grayscale(),
-    transforms.Resize((208, 208)),  # <-- important for matching FC input
+    transforms.Resize((208, 208)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.5], std=[0.5])
 ])
 
-# ✅ Tumor Prediction Function
+# ─── Gradio Prediction Functions ───────────────────────────────────────────────
+
 def predict_tumor(image):
-    image = transform(image).unsqueeze(0)
+    tensor = transform(image).unsqueeze(0)
     with torch.no_grad():
-        out = tumor_model(image)
+        out = tumor_model(tensor)
         pred = torch.argmax(out, dim=1).item()
-        return tumor_labels[pred]
+    return tumor_labels[pred]
 
-# ✅ Glioma Stage Prediction Function
 def predict_stage(gender, age, idh1, tp53, atrx, pten, egfr, cic, pik3ca):
     gender_val = 0 if gender == "Male" else 1
     features = [gender_val, age, idh1, tp53, atrx, pten, egfr, cic, pik3ca]
@@ -42,41 +52,33 @@ def predict_stage(gender, age, idh1, tp53, atrx, pten, egfr, cic, pik3ca):
     with torch.no_grad():
         out = glioma_model(x)
         pred = torch.argmax(out, dim=1).item()
-        return stage_labels[pred]
+    return stage_labels[pred]
+
+# ─── Gradio UI ────────────────────────────────────────────────────────────────
 
-# ✅ Tumor Detection Tab
 tumor_tab = gr.Interface(
     fn=predict_tumor,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(),
-    title="🧠 Brain Tumor Detector",
-    description="Upload an MRI image to classify tumor type: glioma, meningioma, notumor, or pituitary."
+    title="Brain Tumor Detector"
 )
 
-# ✅ Glioma Stage Prediction Tab
 stage_tab = gr.Interface(
     fn=predict_stage,
     inputs=[
         gr.Radio(["Male", "Female"], label="Gender"),
         gr.Slider(0, 100, label="Age"),
-        gr.Slider(0, 1, step=1, label="IDH1 Mutation"),
-        gr.Slider(0, 1, step=1, label="TP53 Mutation"),
-        gr.Slider(0, 1, step=1, label="ATRX Mutation"),
-        gr.Slider(0, 1, step=1, label="PTEN Mutation"),
-        gr.Slider(0, 1, step=1, label="EGFR Mutation"),
-        gr.Slider(0, 1, step=1, label="CIC Mutation"),
-        gr.Slider(0, 1, step=1, label="PIK3CA Mutation")
+        gr.Slider(0, 1, step=1, label="IDH1"),
+        gr.Slider(0, 1, step=1, label="TP53"),
+        gr.Slider(0, 1, step=1, label="ATRX"),
+        gr.Slider(0, 1, step=1, label="PTEN"),
+        gr.Slider(0, 1, step=1, label="EGFR"),
+        gr.Slider(0, 1, step=1, label="CIC"),
+        gr.Slider(0, 1, step=1, label="PIK3CA")
     ],
     outputs=gr.Label(),
-    title="🧬 Glioma Stage Classifier",
-    description="Enter mutation and demographic data to classify glioma stage."
-)
-
-# ✅ Combine into a tabbed interface
-demo = gr.TabbedInterface(
-    [tumor_tab, stage_tab],
-    tab_names=["Tumor Detector", "Glioma Stage Predictor"]
+    title="Glioma Stage Predictor"
 )
 
-# ✅ Launch the app
+demo = gr.TabbedInterface([tumor_tab, stage_tab], tab_names=["Tumor Detector", "Glioma Stage"])
 demo.launch()