Update newapi.py

newapi.py

@@ -28,7 +28,7 @@ app.add_middleware(
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    transforms.Normalize(mean=[0.5], std=[0.5]),
 ])
 
 # Define your model directly inside this file (to avoid import errors)
@@ -37,19 +37,23 @@ import torch.nn as nn
 class BrainTumorModel(nn.Module):
     def __init__(self):
         super(BrainTumorModel, self).__init__()
-        self.model = nn.Sequential(
-            nn.Conv2d(3, 16, kernel_size=3),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-            nn.Conv2d(16, 32, kernel_size=3),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-            nn.Flatten(),
-            nn.Linear(32 * 54 * 54, 2),
-        )
+        self.con1d = nn.Conv2d(3, 32, kernel_size=3)
+        self.con2d = nn.Conv2d(32, 64, kernel_size=3)
+        self.con3d = nn.Conv2d(64, 128, kernel_size=3)
+        self.pool = nn.MaxPool2d(2)
+        self.fc1 = nn.Linear(128 * 25 * 25, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.output = nn.Linear(128, 2)
 
     def forward(self, x):
-        return self.model(x)
+        x = self.pool(torch.relu(self.con1d(x)))
+        x = self.pool(torch.relu(self.con2d(x)))
+        x = self.pool(torch.relu(self.con3d(x)))
+        x = x.view(-1, 128 * 25 * 25)
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.output(x)
+        return x
 
 # Load model
 model_path = "BTD_model.pth"