Update TumorModel.py

TumorModel.py

@@ -1,29 +1,31 @@
 import torch.nn as nn
+import torch
 
 class TumorClassification(nn.Module):
     def __init__(self):
         super().__init__()
-        self.model = nn.Sequential(
-            nn.Conv2d(1, 32, 3, 1, 1),      # con1d
-            nn.ReLU(),
-            nn.MaxPool2d(2),
+        self.con1d = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
+        self.con2d = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+        self.con3d = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+        
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
 
-            nn.Conv2d(32, 64, 3, 1, 1),     # con2d
-            nn.ReLU(),
-            nn.MaxPool2d(2),
+        self.fc1 = nn.Linear(128 * 26 * 26, 512)  # Must match training input size (208x208 image)
+        self.fc2 = nn.Linear(512, 256)
+        self.output = nn.Linear(256, 4)
 
-            nn.Conv2d(64, 128, 3, 1, 1),    # con3d
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-
-            nn.Flatten(),
-            nn.Linear(128 * 26 * 26, 512),  # 128×26×26 = 86528
-            nn.ReLU(),
-            nn.Linear(512, 4)               # 4 classes
-        )
+        self.relu = nn.ReLU()
 
     def forward(self, x):
-        return self.model(x)
+        x = self.pool(self.relu(self.con1d(x)))
+        x = self.pool(self.relu(self.con2d(x)))
+        x = self.pool(self.relu(self.con3d(x)))
+        x = x.view(x.size(0), -1)
+        x = self.relu(self.fc1(x))
+        x = self.relu(self.fc2(x))
+        x = self.output(x)
+        return x
+
 
 class GliomaStageModel(nn.Module):
     def __init__(self):