modify model load process

README.md

@@ -9,12 +9,31 @@ lon_std = 0.0006184167829766685
 # TO RUN: 
 
 from huggingface_hub import hf_hub_download
+import torchvision.models as models
 import torch
+import torch.nn as nn
 
-repo_id = "cis-5190-final-fall24/ImageToGPSproject_model"
+# Specify the repository and the filename of the model you want to load
+repo_id = "cis-5190-final-fall24/ImageToGPSproject_model"  # Replace with your repo name
 filename = "final_model.pth"
 
+class ResNetGPSModel(nn.Module):
+    def __init__(self):
+        super(ResNetGPSModel, self).__init__()
+        self.resnet = models.resnet101()  # Updated for PyTorch >=0.13
+        self.resnet.fc = nn.Sequential(
+            nn.Dropout(0.4),  # Dropout for regularization
+            nn.Linear(self.resnet.fc.in_features, 2)  # Latitude and Longitude
+        )
+
+    def forward(self, x):
+        return self.resnet(x)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = ResNetGPSModel().to(device)
 model_path = hf_hub_download(repo_id=repo_id, filename=filename)
 
-model_test = torch.load(model_path)
-model_test.eval()
+# Load the model using torch
+state_dict = torch.load(model_path)
+model.load_state_dict(state_dict)
+model.eval()  # Set the model to evaluation mode