Codewithsalty commited on
Commit
81484f1
·
verified ·
1 Parent(s): 0b9429c

Update TumorModel.py

Browse files
Files changed (1) hide show
  1. TumorModel.py +23 -1
TumorModel.py CHANGED
@@ -1,10 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  class GliomaStageModel(nn.Module):
2
  def __init__(self):
3
  super().__init__()
4
  self.fc1 = nn.Linear(9, 100)
5
  self.fc2 = nn.Linear(100, 50)
6
  self.fc3 = nn.Linear(50, 30)
7
- self.out = nn.Linear(30, 2) # ✅ MATCHING saved model weights
8
 
9
  def forward(self, x):
10
  x = nn.functional.relu(self.fc1(x))
 
1
+ import torch.nn as nn
2
+
3
+ class TumorClassification(nn.Module):
4
+ def __init__(self):
5
+ super().__init__()
6
+ self.con1d = nn.Conv2d(1, 32, kernel_size=3, padding=1)
7
+ self.con2d = nn.Conv2d(32, 64, kernel_size=3, padding=1)
8
+ self.con3d = nn.Conv2d(64, 128, kernel_size=3, padding=1)
9
+ self.pool = nn.MaxPool2d(2, 2)
10
+ self.fc1 = nn.Linear(128 * 28 * 28, 512)
11
+ self.fc2 = nn.Linear(512, 256)
12
+ self.output = nn.Linear(256, 4)
13
+
14
+ def forward(self, x):
15
+ x = nn.functional.relu(self.pool(self.con1d(x)))
16
+ x = nn.functional.relu(self.pool(self.con2d(x)))
17
+ x = nn.functional.relu(self.pool(self.con3d(x)))
18
+ x = x.view(x.size(0), -1)
19
+ x = nn.functional.relu(self.fc1(x))
20
+ x = nn.functional.relu(self.fc2(x))
21
+ return self.output(x)
22
+
23
  class GliomaStageModel(nn.Module):
24
  def __init__(self):
25
  super().__init__()
26
  self.fc1 = nn.Linear(9, 100)
27
  self.fc2 = nn.Linear(100, 50)
28
  self.fc3 = nn.Linear(50, 30)
29
+ self.out = nn.Linear(30, 2) # ✅ Make sure this matches your trained model
30
 
31
  def forward(self, x):
32
  x = nn.functional.relu(self.fc1(x))