IanAndJohn/Model_Ian

latitude_mean: 39.951631102585964
latitude_std: 0.0006960598068888123
longitude_mean: -75.1914340210287
longitude_std: 0.0006455062924978866

from huggingface_hub import hf_hub_download
import torch

import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import PyTorchModelHubMixin
import torchvision.models as models

class SimpleCNN(nn.Module, PyTorchModelHubMixin):
    def __init__(self):
        super().__init__()

        # Convolutional layers
        self.conv3to32 = nn.Conv2d(in_channels=3, out_channels=15, kernel_size=9, stride=1, padding=4)

        self.conv32to32kernel5 = nn.Conv2d(in_channels=15, out_channels=15, kernel_size=5, stride=1, padding=2)

        self.conv32to64 = nn.Conv2d(in_channels=15, out_channels=30, kernel_size=3, stride=1, padding=1)

        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.dropout = nn.Dropout(0.5)

        self.linear_input_dims = 30*56*56
        self.fc_1 = nn.Linear(self.linear_input_dims, 100)
        self.fc_2 = nn.Linear(100, 2)

    def forward(self, x):
        x = F.relu(self.conv3to32(x))
        x = F.relu(self.conv32to32kernel5(x))


        x = self.pool2(x)
        x = F.relu(self.conv32to64(x))

        x = self.pool2(x)
        x = self.dropout(x)

        x = x.view(-1, self.linear_input_dims)
        x = F.relu(self.fc_1(x))
        x = self.fc_2(x)
        return x

    def save_model(self, save_path):
        """Save model locally using the Hugging Face format."""
        self.save_pretrained(save_path)

    def push_model(self, repo_name):
        """Push the model to the Hugging Face Hub."""
        self.push_to_hub(repo_name)


# Specify the repository and the filename of the model you want to load
repo_id = "IanAndJohn/Model_Ian"  # Replace with your repo name
filename = model_save_path

model_path = hf_hub_download(repo_id=repo_id, filename=filename)

# Load the model using torch
model = SimpleCNN()
model.load_state_dict(torch.load(model_path))
model.eval()  # Set the model to evaluation mode