Update app.py

app.py

@@ -1,79 +1,79 @@
-import torch
-import gradio as gr
-from PIL import Image
-from torchvision import transforms
-from TumorModel import TumorClassification, GliomaStageModel
-
-# 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
-glioma_model = GliomaStageModel()
-glioma_model.load_state_dict(torch.load("glioma_stages.pth", map_location=torch.device("cpu")))
-glioma_model.eval()
-
-# Labels
-tumor_labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
-stage_labels = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
-
-# Image transform
-transform = transforms.Compose([
-    transforms.Grayscale(),
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5], std=[0.5])
-])
-
-# Tumor type prediction
-def predict_tumor(image):
-    image = transform(image).unsqueeze(0)
-    with torch.no_grad():
-        out = tumor_model(image)
-        pred = torch.argmax(out, dim=1).item()
-        return tumor_labels[pred]
-
-# Glioma stage prediction
-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]
-    x = torch.tensor(features).float().unsqueeze(0)
-    with torch.no_grad():
-        out = glioma_model(x)
-        pred = torch.argmax(out, dim=1).item()
-        return stage_labels[pred]
-
-# Interface 1: Tumor Detector
-tumor_tab = gr.Interface(
-    fn=predict_tumor,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(),
-    title="🧠 Brain Tumor Detection",
-    description="Upload an MRI image to classify tumor type: glioma, meningioma, notumor, or pituitary."
-)
-
-# Interface 2: Glioma Stage Prediction
-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")
-    ],
-    outputs=gr.Label(),
-    title="🧬 Glioma Stage Predictor",
-    description="Enter patient’s mutation and demographic data to predict glioma stage."
-)
-
-# Combine in tabs
-demo = gr.TabbedInterface([tumor_tab, stage_tab], tab_names=["Tumor Detector", "Glioma Stage"])
-
-# Run app
-demo.launch()
+import torch
+import gradio as gr
+from PIL import Image
+from torchvision import transforms
+from TumorModel import TumorClassification, GliomaStageModel
+
+# 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
+glioma_model = GliomaStageModel()
+glioma_model.load_state_dict(torch.load("glioma_stages.pth", map_location=torch.device("cpu")))
+glioma_model.eval()
+
+# Labels
+tumor_labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
+stage_labels = ['Stage 1', 'Stage 2']  # Only 2 classes in your model
+
+# Transform for image input
+transform = transforms.Compose([
+    transforms.Grayscale(),
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5], std=[0.5])
+])
+
+# Predict tumor type
+def predict_tumor(image):
+    image = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        out = tumor_model(image)
+        pred = torch.argmax(out, dim=1).item()
+        return tumor_labels[pred]
+
+# Predict glioma stage
+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]
+    x = torch.tensor(features).float().unsqueeze(0)
+    with torch.no_grad():
+        out = glioma_model(x)
+        pred = torch.argmax(out, dim=1).item()
+        return stage_labels[pred]
+
+# Interface 1: Tumor Classification
+tumor_tab = gr.Interface(
+    fn=predict_tumor,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(),
+    title="🧠 Brain Tumor Detection",
+    description="Upload an MRI image to classify tumor type: glioma, meningioma, notumor, or pituitary."
+)
+
+# Interface 2: Glioma Stage Prediction
+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")
+    ],
+    outputs=gr.Label(),
+    title="🧬 Glioma Stage Classifier",
+    description="Enter patient 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"])
+
+# Launch
+demo.launch()