Update README.md

README.md

@@ -10,44 +10,40 @@ 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-resnet_vit-base"  # Replace with your repo name
-filename = "resnet_vit_gps_regressor_complete.pth"
+repo_id = "FinalProj5190/vit_base_72"  # Replace with your repo name
+filename = "resnet_gps_regressor_complete.pth"
 
 model_path = hf_hub_download(repo_id=repo_id, filename=filename)
 
-# Load the model using torch
-model_test = torch.load(model_path)
-model_test.eval()  # Set the model to evaluation mode
+model_test = MultiModalModel()
+model_test.load_state_dict(torch.load(model_path))
+model_test.eval()
 ```
 
 The model implementation is here:
 ```
-from transformers import ViTModel
-class HybridGPSModel(nn.Module):
-    def __init__(self, num_classes=2):
-        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')
-
-        # Combined regression head
-        self.regression_head = nn.Sequential(
-            nn.Linear(512 + self.vit.config.hidden_size, 128),
+from transformers import AutoModel
+
+class MultiModalModel(nn.Module):
+    def __init__(self, image_model_name='google/vit-base-patch16-224-in21k', output_dim=2):
+        super(MultiModalModel, self).__init__()
+
+        # Load Vision Transformer for feature extraction
+        self.image_model = AutoModel.from_pretrained(image_model_name, output_hidden_states=True)
+
+        # Combine image and GPS features for regression
+        self.regressor = nn.Sequential(
+            nn.Linear(self.image_model.config.hidden_size, 128),
             nn.ReLU(),
-            nn.Linear(128, num_classes) 
+            nn.Dropout(0.3),
+            nn.Linear(128, output_dim),
         )
 
-    def forward(self, x):
-        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)
+    def forward(self, image):
+        # Extract image features from the last hidden state
+        image_outputs = self.image_model(image)
+        image_features = image_outputs.last_hidden_state[:, 0, :]  # CLS token features
 
-        # Predict GPS coordinates
-        gps_coordinates = self.regression_head(combined_features)
-        return gps_coordinates
+        # Final regression
+        return self.regressor(image_features)
 ```