Update model.py

model.py

@@ -7,21 +7,23 @@ class SimpleCNN(nn.Module):
         super(SimpleCNN, self).__init__()
         self.num_classes = num_classes
         self.model_type = model_type
-        # Define convolutional and fc layers based on model_type
+        # Model architectures assume 224x224 input
         if model_type == 'f':
-            # Two pool layers: 448 -> 224 -> 112 -> 56 -> 28
+            # After 3 pool layers: 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':
+            # After 3 pool layers: 224 -> 112 -> 56 -> 28
             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':
+            # After 4 pool layers: 224 -> 112 -> 56 -> 28 -> 14
             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)
@@ -51,6 +53,7 @@ class SimpleCNN(nn.Module):
 def load_model(version='c', device='cpu'):
     """
     Downloads and loads the SimpleCNN model for the specified version: 'f', 'c', or 'q'.
+    Input images must be resized to 224x224.
     """
     model_type = version.lower()
     filename = f"Vbai-2.1{model_type}.pt"
@@ -65,6 +68,6 @@ def load_model(version='c', device='cpu'):
     # Initialize and load model
     model = SimpleCNN(model_type=model_type, num_classes=6).to(device)
     state_dict = torch.load(weights_path, map_location=device)
-    model.load_state_dict(state_dict)
+    model.load_state_dict(state_dict, strict=False)
     model.eval()
     return model