Update README.md

README.md

@@ -10,8 +10,8 @@ from huggingface_hub import hf_hub_download
 import torch
 
 # Specify the repository and the filename of the model you want to load
-repo_id = "FinalProj5190/ImageToGPSproject_new_vit"  # Replace with your repo name
-filename = "resnet_gps_regressor_complete.pth"
+repo_id = "FinalProj5190/ImageToGPSproject-resnet_vit-base"  # Replace with your repo name
+filename = "resnet_vit_gps_regressor_complete.pth"
 
 model_path = hf_hub_download(repo_id=repo_id, filename=filename)
 
@@ -23,23 +23,31 @@ model_test.eval()  # Set the model to evaluation mode
 The model implementation is here:
 ```
 from transformers import ViTModel
-class MultiModalModel(nn.Module):
+class HybridGPSModel(nn.Module):
     def __init__(self, num_classes=2):
-        super(MultiModalModel, self).__init__()
+        super(HybridGPSModel, self).__init__()
+        # Pre-trained ResNet for feature extraction
+        self.resnet = resnet18(pretrained=True)
+        self.resnet.fc = nn.Identity()
+
+        # Pre-trained Vision Transformer
         self.vit = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
-        
-        # Replace for regression instead of classification
+
+        # Combined regression head
         self.regression_head = nn.Sequential(
-            nn.Linear(self.vit.config.hidden_size, 512),
+            nn.Linear(512 + self.vit.config.hidden_size, 128),
             nn.ReLU(),
-            nn.Linear(512, num_classes)
+            nn.Linear(128, num_classes) 
         )
-    
+
     def forward(self, x):
-        outputs = self.vit(pixel_values=x)
-        # Take the last hidden state (CLS token embedding)
-        cls_output = outputs.last_hidden_state[:, 0, :]
-        # Pass through the regression head
-        gps_coordinates = self.regression_head(cls_output)
+        resnet_features = self.resnet(x)
+        vit_outputs = self.vit(pixel_values=x)
+        vit_features = vit_outputs.last_hidden_state[:, 0, :]  # CLS token
+
+        combined_features = torch.cat((resnet_features, vit_features), dim=1)
+
+        # Predict GPS coordinates
+        gps_coordinates = self.regression_head(combined_features)
         return gps_coordinates
 ```