Update app.py

app.py

@@ -4,29 +4,29 @@ 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")))
 tumor_model.eval()
 
-# Load glioma stage model
+# ✅ Load glioma stage classification model
 glioma_model = GliomaStageModel()
 glioma_model.load_state_dict(torch.load("glioma_stages.pth", map_location=torch.device("cpu")))
 glioma_model.eval()
 
-# Labels
+# ✅ Labels
 tumor_labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
-stage_labels = ['Stage 1', 'Stage 2']  # Only 2 classes in your model
+stage_labels = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
 
-# Transform for image input
+# ✅ Transform (resize to 208x208 to match training)
 transform = transforms.Compose([
     transforms.Grayscale(),
-    transforms.Resize((224, 224)),
+    transforms.Resize((208, 208)),  # <-- important for matching FC input
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.5], std=[0.5])
 ])
 
-# Predict tumor type
+# ✅ Tumor Prediction Function
 def predict_tumor(image):
     image = transform(image).unsqueeze(0)
     with torch.no_grad():
@@ -34,7 +34,7 @@ def predict_tumor(image):
         pred = torch.argmax(out, dim=1).item()
         return tumor_labels[pred]
 
-# Predict glioma stage
+# ✅ 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]
@@ -44,16 +44,16 @@ def predict_stage(gender, age, idh1, tp53, atrx, pten, egfr, cic, pik3ca):
         pred = torch.argmax(out, dim=1).item()
         return stage_labels[pred]
 
-# Interface 1: Tumor Classification
+# ✅ Tumor Detection Tab
 tumor_tab = gr.Interface(
     fn=predict_tumor,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(),
-    title="🧠 Brain Tumor Detection",
+    title="🧠 Brain Tumor Detector",
     description="Upload an MRI image to classify tumor type: glioma, meningioma, notumor, or pituitary."
 )
 
-# Interface 2: Glioma Stage Prediction
+# ✅ Glioma Stage Prediction Tab
 stage_tab = gr.Interface(
     fn=predict_stage,
     inputs=[
@@ -69,11 +69,14 @@ stage_tab = gr.Interface(
     ],
     outputs=gr.Label(),
     title="🧬 Glioma Stage Classifier",
-    description="Enter patient mutation and demographic data to classify glioma stage."
+    description="Enter mutation and demographic data to classify glioma stage."
 )
 
-# Combine both into a tabbed interface
-demo = gr.TabbedInterface([tumor_tab, stage_tab], tab_names=["Tumor Detector", "Glioma Stage"])
+# ✅ Combine into a tabbed interface
+demo = gr.TabbedInterface(
+    [tumor_tab, stage_tab],
+    tab_names=["Tumor Detector", "Glioma Stage Predictor"]
+)
 
-# Launch
+# ✅ Launch the app
 demo.launch()