Update TumorModel.py

TumorModel.py

@@ -1,5 +1,6 @@
 import torch.nn as nn
 
+# 🎯 Tumor Type Classification Model
 class TumorClassification(nn.Module):
     def __init__(self):
         super(TumorClassification, self).__init__()
@@ -16,9 +17,7 @@ class TumorClassification(nn.Module):
         self.pool3 = nn.MaxPool2d(2)
 
         self.flatten = nn.Flatten()
-
-        # 👇 Update this line to match the saved .pth model
-        self.fc1 = nn.Linear(86528, 512)  # Fixed input size
+        self.fc1 = nn.Linear(86528, 512)  # fixed to match pth
         self.relu_fc = nn.ReLU()
         self.fc2 = nn.Linear(512, 256)
         self.relu_fc2 = nn.ReLU()
@@ -31,5 +30,19 @@ class TumorClassification(nn.Module):
         x = self.flatten(x)
         x = self.relu_fc(self.fc1(x))
         x = self.relu_fc2(self.fc2(x))
-        x = self.output(x)
-        return x
+        return self.output(x)
+
+# 🧬 Glioma Stage Classifier Model
+class GliomaStageModel(nn.Module):
+    def __init__(self):
+        super(GliomaStageModel, self).__init__()
+        self.model = nn.Sequential(
+            nn.Linear(9, 128),
+            nn.ReLU(),
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Linear(64, 4)
+        )
+
+    def forward(self, x):
+        return self.model(x)