Views
No views yet
1# Import all the dependencies
2import torch
3import torch.nn as nn
4import torchvision.models as models
5import torchvision.transforms as transforms
6from torch.utils.data import DataLoader, Dataset
7from transformers import AutoImageProcessor, AutoModelForImageClassification, AutoModel
8from huggingface_hub import PyTorchModelHubMixin
9from PIL import Image
10import os
11import numpy as np
12
13
14class EfficientNetGPSModel(nn.Module, PyTorchModelHubMixin):
15 def __init__(self, eff_name="efficientnet_b0", num_outputs=2):
16 super(EfficientNetGPSModel, self).__init__()
17
18 # Load the EfficientNet backbone
19 self.efficientnet = getattr(models, eff_name)(pretrained=True)
20
21 # Replace the classifier head while keeping the overall structure simple
22 in_features = self.efficientnet.classifier[1].in_features
23 self.efficientnet.classifier = nn.Sequential(
24 nn.Linear(in_features, num_outputs) # Directly map to GPS coordinates
25 )
26
27 def forward(self, x):
28 return self.efficientnet(x)
29
30 def save_model(self, save_path):
31 self.save_pretrained(save_path)
32
33 def push_model(self, repo_name):
34 self.push_to_hub(repo_name)model = EfficientNetGPSModel.from_pretrained("cis519/efficient-Net")