ImageToGPSproject_convnext_mobilenet

Image to GPS Project - ConvNext and MobileNet Ensemble

## Training Data Statistics
lat_mean = 39.951537011424264
lat_std = 0.0006940325318781937
lon_mean = -75.19152009539549
lon_std = 0.0007607716964655242

How to Load the Model and Perform Inference

# install dependencies
pip install geopy datasets torch torchvision huggingface_hub

# import packages
import numpy as np
from geopy.distance import geodesic
import torch
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
import torch.nn as nn
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights, convnext_tiny, ConvNeXt_Tiny_Weights
from datasets import load_dataset
from huggingface_hub import hf_hub_download

# load the model 
repo_id = "cis519projectA/ImageToGPSproject_convnext_mobilenet"
filename = "convnext_mobilenet_ensemble_model.pth"
model_path = hf_hub_download(repo_id=repo_id, filename=filename)

# define models
class CustomConvNeXtModel(nn.Module):
    def __init__(self, weights=ConvNeXt_Tiny_Weights.DEFAULT, num_classes=2):
        super().__init__()

        # Load pre-trained ConvNeXt model
        self.convnext = convnext_tiny(weights=weights)
        in_features = self.convnext.classifier[2].in_features

        self.convnext.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(in_features, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(p=0.3),
            nn.Linear(512, num_classes)
        )

        # Freeze early layers
        for param in self.convnext.features[:4].parameters():
            param.requires_grad = False

    def forward(self, x):
        return self.convnext(x)

class CustomMobileNetModel(nn.Module):
    def __init__(self, weights=MobileNet_V2_Weights.DEFAULT, num_classes=2):
        super().__init__()

        # Load pre-trained MobileNet model
        self.mobilenet = mobilenet_v2(weights=weights)
        in_features = self.mobilenet.classifier[1].in_features

        self.mobilenet.classifier = nn.Sequential(
            nn.Linear(in_features, 1024),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(512, num_classes)
        )

        # Freeze early layers
        for param in self.mobilenet.features[:5].parameters():
            param.requires_grad = False

    def forward(self, x):
        return self.mobilenet(x)

class EnsembleModel(nn.Module):
    def __init__(self, resnet_model, mobilenet_model, num_classes=2):
        super().__init__()
        self.resnet = resnet_model
        self.mobilenet = mobilenet_model
        self.fc = nn.Sequential(
            nn.Linear(num_classes * 2, 512), 
            nn.ReLU(),
            nn.Dropout(p=0.3),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        resnet_out = self.resnet(x)
        mobilenet_out = self.mobilenet(x)
        combined = torch.cat((resnet_out, mobilenet_out), dim=1)
        output = self.fc(combined)
        return output

convnext_model = CustomConvNeXtModel()
mobilenet_model = CustomMobileNetModel()
ensemble_model = EnsembleModel(convnext_model, mobilenet_model)

# load the model weights
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
state_dict = torch.load(model_path, map_location=device)
ensemble_model.load_state_dict(state_dict)
ensemble_model.to(device)
ensemble_model.eval()


# load the dataset
dataset_test = load_dataset("gydou/released_img", split="train")

# define transformers
inference_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Parameters for denormalization
lat_mean = 39.951537011424264
lat_std = 0.0006940325318781937
lon_mean = -75.19152009539549
lon_std = 0.0007607716964655242

class GPSImageDataset(Dataset):
    def __init__(self, hf_dataset, transform=None, lat_mean=None, lat_std=None, lon_mean=None, lon_std=None):
        self.hf_dataset = hf_dataset
        self.transform = transform
        self.latitude_mean = lat_mean
        self.latitude_std = lat_std
        self.longitude_mean = lon_mean
        self.longitude_std = lon_std

    def __len__(self):
        return len(self.hf_dataset)

    def __getitem__(self, idx):
        example = self.hf_dataset[idx]
        image = example['image']
        latitude = example['Latitude']
        longitude = example['Longitude']
        if self.transform:
            image = self.transform(image)
        latitude = (latitude - self.latitude_mean) / self.latitude_std
        longitude = (longitude - self.longitude_mean) / self.longitude_std
        gps_coords = torch.tensor([latitude, longitude], dtype=torch.float32)
        return image, gps_coords

# transform test data
test_dataset = GPSImageDataset(
    hf_dataset=dataset_test,
    transform=inference_transform,
    lat_mean=lat_mean,
    lat_std=lat_std,
    lon_mean=lon_mean,
    lon_std=lon_std
)
test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

# evaluate
def evaluate_model_single_batch(model, dataloader, lat_mean, lat_std, lon_mean, lon_std):
    all_distances = []
    model.eval()
    with torch.no_grad():
        for batch_idx, (images, gps_coords) in enumerate(dataloader):            
            images, gps_coords = images.to(device), gps_coords.to(device)
            outputs = model(images)
            preds_denorm = outputs.cpu().numpy() * np.array([lat_std, lon_std]) + np.array([lat_mean, lon_mean])
            actuals_denorm = gps_coords.cpu().numpy() * np.array([lat_std, lon_std]) + np.array([lat_mean, lon_mean])
            for pred, actual in zip(preds_denorm, actuals_denorm):
                distance = geodesic((actual[0], actual[1]), (pred[0], pred[1])).meters
                all_distances.append(distance)
            break

    mean_error = np.mean(all_distances)
    rmse_error = np.sqrt(np.mean(np.square(all_distances)))
    return mean_error, rmse_error


# Evaluate using only one batch
mean_error, rmse_error = evaluate_model_single_batch(
    ensemble_model, test_dataloader, lat_mean, lat_std, lon_mean, lon_std
)
print(f"Mean Error (meters): {mean_error:.2f}, RMSE (meters): {rmse_error:.2f}")