Add model.py for project submission

model.py

@@ -1,46 +1,50 @@
 
 import torch
 import torch.nn as nn
-import timm
+import torchvision.models as models
 
 class IMG2GPS(nn.Module):
     """
     EfficientNet-B0 model for GPS coordinate prediction from images.
-
+    
     Input: Batch of images (N, 3, 224, 224) - ImageNet normalized
     Output: Batch of GPS coordinates (N, 2) - raw lat/lon in degrees
     """
-
+    
     def __init__(self):
         super().__init__()
-
-        # Load pre-trained EfficientNet-B0
-        self.backbone = timm.create_model('efficientnet_b0', pretrained=False, num_classes=2)
-
+        
+        # Load pre-trained EfficientNet-B0 from torchvision
+        self.backbone = models.efficientnet_b0(pretrained=False)
+        
+        # Replace the final classifier layer to output 2 values
+        num_features = self.backbone.classifier[1].in_features
+        self.backbone.classifier[1] = nn.Linear(num_features, 2)
+        
         # Hardcoded normalization statistics from training set
         self.lat_mean = 39.951525
         self.lat_std = 0.000652
         self.lon_mean = -75.191400
         self.lon_std = 0.000598
-
+    
     def forward(self, x):
         """
         Forward pass through the model.
-
+        
         Args:
             x: Input tensor of shape (N, 3, 224, 224) - normalized images
                or list of tensors
-
+        
         Returns:
             Tensor of shape (N, 2) - denormalized lat/lon in degrees
         """
         # Handle case where input is a list of tensors
         if isinstance(x, list):
             x = torch.stack(x)
-
+        
         # Model outputs normalized GPS coordinates
         normalized_coords = self.backbone(x)  # Shape: (N, 2)
-
+        
         # Denormalize to get raw lat/lon in degrees
         denormalized_coords = normalized_coords * torch.tensor(
             [self.lat_std, self.lon_std],
@@ -51,24 +55,24 @@ class IMG2GPS(nn.Module):
             device=x.device,
             dtype=x.dtype
         )
-
+        
         return denormalized_coords
-
+    
     def predict(self, batch):
         """
         Inference method for compatibility with backend.
-
+        
         Args:
             batch: Input tensor of shape (N, 3, 224, 224)
                    or list of tensors
-
+        
         Returns:
             numpy array of shape (N, 2) with raw lat/lon in degrees
         """
         # Handle case where batch is a list of tensors
         if isinstance(batch, list):
             batch = torch.stack(batch)
-
+        
         self.eval()
         with torch.no_grad():
             output = self.forward(batch)
@@ -78,7 +82,7 @@ class IMG2GPS(nn.Module):
 def get_model():
     """
     Factory function to instantiate the model.
-
+    
     Returns:
         IMG2GPS model instance
     """