Update model.py

model.py

@@ -7,23 +7,19 @@ class SimpleCNN(nn.Module):
         super(SimpleCNN, self).__init__()
         self.num_classes = num_classes
         self.model_type = model_type
-        # Model architectures assume 224x224 input
         if model_type == 'f':
-            # 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,21 +47,15 @@ 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"
 
-    # Download the weight file from Hugging Face Hub
     weights_path = hf_hub_download(
         repo_id="Neurazum/Vbai-DPA-2.1",
         filename=filename,
         repo_type="model"
     )
 
-    # 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, strict=False)