Update app.py

app.py

@@ -7,16 +7,20 @@ import torchvision.transforms as transforms
 from PIL import Image
 import matplotlib.pyplot as plt
 
-# Define model and load state
-class GeoNet(nn.Module):
-    # Define your model architecture here
+class geoNet(nn.Module):
     def __init__(self):
-        super(GeoNet, self).__init__()
-        # Define layers here
+        super(geoNet, self).__init__()
+        self.name = "geo"
+        self.fc1 = nn.Linear(256*6*6, 244)
+        self.classifier = nn.Linear(244, 63) # Classification head
+        self.regressor = nn.Linear(244, 2) # Regression head for (latitude, longitude)
 
     def forward(self, x):
-        # Define forward pass here
-        return x
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.fc1(x))
+        province_pred = self.classifier(x)
+        coords_pred = self.regressor(x)
+        return province_pred, coords_pred
 
 model = GeoNet()
 model_path = 'model_geo_epoch79.pth'