Update model.py

model.py

@@ -5,59 +5,57 @@ from huggingface_hub import hf_hub_download
 class SimpleCNN(nn.Module):
     def __init__(self, model_type='f', num_classes=6):
         super(SimpleCNN, self).__init__()
+        self.num_classes = num_classes
         self.model_type = model_type
-        # Define conv layers based on model_type
+        # Define convolutional and fc layers based on model_type
         if model_type == 'f':
-            channels = [3, 16, 32, 64]
-            dropout_p = 0.5
-            fc_hidden = 256
+            # Two pool layers: 448 -> 224 -> 112 -> 56 -> 28
+            self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
+            self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
+            self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+            self.fc1 = nn.Linear(64 * 28 * 28, 256)
+            self.dropout = nn.Dropout(0.5)
         elif model_type == 'c':
-            channels = [3, 32, 64, 128]
-            dropout_p = 0.5
-            fc_hidden = 512
+            self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
+            self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+            self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+            self.fc1 = nn.Linear(128 * 28 * 28, 512)
+            self.dropout = nn.Dropout(0.5)
         elif model_type == 'q':
-            channels = [3, 64, 128, 256, 512]
-            dropout_p = 0.3
-            fc_hidden = 1024
+            self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
+            self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+            self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
+            self.conv4 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
+            self.fc1 = nn.Linear(512 * 14 * 14, 1024)
+            self.dropout = nn.Dropout(0.3)
         else:
             raise ValueError(f"Unknown model type: {model_type}")
 
-        # Build conv blocks
-        layers = []
-        for in_c, out_c in zip(channels[:-1], channels[1:]):
-            layers.append(nn.Conv2d(in_c, out_c, kernel_size=3, padding=1))
-            layers.append(nn.ReLU())
-            layers.append(nn.MaxPool2d(2))
-        self.features = nn.Sequential(*layers)
-        self.dropout = nn.Dropout(dropout_p)
-
-        # Dynamically compute flattened size
-        with torch.no_grad():
-            dummy = torch.zeros(1, 3, 448, 448)
-            feat = self.features(dummy)
-            flattened_size = feat.view(1, -1).size(1)
-
-        # Fully connected layers
-        self.fc1 = nn.Linear(flattened_size, fc_hidden)
-        self.fc2 = nn.Linear(fc_hidden, num_classes)
+        self.relu = nn.ReLU()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.fc2 = nn.Linear(self.fc1.out_features, num_classes)
 
     def forward(self, x):
-        x = self.features(x)
+        x = self.pool(self.relu(self.conv1(x)))
+        x = self.pool(self.relu(self.conv2(x)))
+        x = self.pool(self.relu(self.conv3(x)))
+        if self.model_type == 'q':
+            x = self.pool(self.relu(self.conv4(x)))
         x = x.view(x.size(0), -1)
-        x = self.dropout(torch.relu(self.fc1(x)))
+        x = self.relu(self.fc1(x))
+        x = self.dropout(x)
         x = self.fc2(x)
         return x
 
 
 def load_model(version='c', device='cpu'):
     """
-    Loads the correct model based on version: 'f', 'c', or 'q'.
+    Downloads and loads the SimpleCNN model for the specified version: 'f', 'c', or 'q'.
     """
-    # Determine filename and model_type
     model_type = version.lower()
     filename = f"Vbai-2.1{model_type}.pt"
 
-    # Download weights
+    # Download the weight file from Hugging Face Hub
     weights_path = hf_hub_download(
         repo_id="Neurazum/Vbai-DPA-2.1",
         filename=filename,
@@ -66,7 +64,7 @@ def load_model(version='c', device='cpu'):
 
     # Initialize and load model
     model = SimpleCNN(model_type=model_type, num_classes=6).to(device)
-    state = torch.load(weights_path, map_location=device)
-    model.load_state_dict(state)
+    state_dict = torch.load(weights_path, map_location=device)
+    model.load_state_dict(state_dict)
     model.eval()
-    return model
+    return model