Update app.py

app.py

@@ -8,17 +8,19 @@ import gradio as gr
 import plotly.graph_objects as go
 from io import BytesIO
 from PIL import Image
-from torchvision import transforms,models
-from sklearn.preprocessing import LabelEncoder,MinMaxScaler
+from torchvision import transforms, models
+from sklearn.preprocessing import LabelEncoder, MinMaxScaler
 from gradio import Interface, Image, Label, HTML
 from huggingface_hub import snapshot_download
+import s2sphere
+import folium
 
 # Retrieve the token from the environment variables
 token = os.environ.get("token")
 
 # Download the repository snapshot
 local_dir = snapshot_download(
-    repo_id="robocan/GeoG_coordinate",
+    repo_id="robocan/GeoG-GCP",
     repo_type="model",
     local_dir="SVD",
     token=token
@@ -27,7 +29,6 @@ local_dir = snapshot_download(
 device = 'cpu'
 le = LabelEncoder()
 le = joblib.load("SVD/le.gz")
-MMS = joblib.load("SVD/MMS.gz")
 len_classes = len(le.classes_) + 1
 
 class ModelPre(torch.nn.Module):
@@ -36,9 +37,9 @@ class ModelPre(torch.nn.Module):
         self.embedding = torch.nn.Sequential(
             *list(models.convnext_small(weights=models.ConvNeXt_Small_Weights.IMAGENET1K_V1).children())[:-1],
             torch.nn.Flatten(),
-            torch.nn.Linear(in_features=768,out_features=512),
+            torch.nn.Linear(in_features=768, out_features=512),
             torch.nn.ReLU(),
-            torch.nn.Linear(in_features=512,out_features=len_classes),
+            torch.nn.Linear(in_features=512, out_features=len_classes),
         )
         # Freeze all layers
 
@@ -47,30 +48,8 @@ class ModelPre(torch.nn.Module):
     
 # Load the pretrained model 
 model = ModelPre()
-#for param in model.parameters():
-#   param.requires_grad = False
-class GeoGcord(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.embedding = torch.nn.Sequential(
-            *list(model.children())[0][:-1],
-            torch.nn.Linear(in_features=512,out_features=256),
-            torch.nn.ReLU(),
-            torch.nn.Linear(in_features=256,out_features=128),
-            torch.nn.ReLU(),
-            torch.nn.Linear(in_features=128,out_features=2),
-        )
-        # Freeze all layers
 
-    def forward(self, data):
-        return self.embedding(data)
-    
-    
-
-# Load the pre-trained model 
-model = GeoGcord()
 model_w = torch.load("SVD/GeoG.pth", map_location=torch.device(device))
-
 model.load_state_dict(model_w['model'])
 
 cmp = transforms.Compose([
@@ -79,27 +58,45 @@ cmp = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Predict function for the new regression model
 def predict(input_img):
     with torch.inference_mode():
         img = cmp(input_img).unsqueeze(0)
         res = model(img.to(device))
-        # Assuming res is a 2-layer regression output, and MMS.inverse_transform is needed
-        prediction = MMS.inverse_transform(res.cpu().numpy()).flatten()
-        return prediction
+        probabilities = torch.softmax(res, dim=1).cpu().numpy().flatten()
+        top_10_indices = np.argsort(probabilities)[-10:][::-1]
+        top_10_probabilities = probabilities[top_10_indices]
+        top_10_predictions = le.inverse_transform(top_10_indices)
+        
+        results = {top_10_predictions[i]: float(top_10_probabilities[i]) for i in range(10)}
+        return results, top_10_predictions
+
+# Function to get S2 cell polygon
+def get_s2_cell_polygon(cell_id):
+    cell = s2sphere.Cell(s2sphere.CellId(cell_id))
+    vertices = []
+    for i in range(4):
+        vertex = s2sphere.LatLng.from_point(cell.get_vertex(i))
+        vertices.append((vertex.lat().degrees, vertex.lng().degrees))
+    vertices.append(vertices[0])  # Close the polygon
+    return vertices
 
 # Function to generate Plotly map figure
-def create_map_figure(lat, lon):
-    fig = go.Figure(go.Scattermapbox(
-        lat=[lat],
-        lon=[lon],
-        mode='markers',
-        marker=go.scattermapbox.Marker(
-            size=14
-        ),
-        text=[f'Lat: {lat}, Lon: {lon}'],
-        hoverinfo='text'
-    ))
+def create_map_figure(predictions, cell_ids):
+    fig = go.Figure()
+
+    for cell_id in cell_ids:
+        cell_id = int(cell_id)
+        polygon = get_s2_cell_polygon(cell_id)
+        lats, lons = zip(*polygon)
+        fig.add_trace(go.Scattermapbox(
+            lat=lats,
+            lon=lons,
+            mode='lines',
+            fill='toself',
+            fillcolor='rgba(0, 0, 255, 0.2)',
+            line=dict(color='blue'),
+            name=f'Cell ID: {cell_id}'
+        ))
 
     fig.update_layout(
         mapbox_style="open-street-map",
@@ -107,8 +104,8 @@ def create_map_figure(lat, lon):
         mapbox=dict(
             bearing=0,
             center=go.layout.mapbox.Center(
-                lat=lat,
-                lon=lon
+                lat=np.mean(lats),
+                lon=np.mean(lons)
             ),
             pitch=0,
             zoom=3
@@ -119,8 +116,8 @@ def create_map_figure(lat, lon):
 
 # Create label output function
 def create_label_output(predictions):
-    lat, lon = predictions
-    fig = create_map_figure(lat, lon)
+    results, cell_ids = predictions
+    fig = create_map_figure(results, cell_ids)
     return fig
 
 # Predict and plot function
@@ -138,4 +135,4 @@ with gr.Blocks() as gradio_app:
         btn_predict.click(predict_and_plot, inputs=input_image, outputs=output_map)
         examples = ["GB.PNG", "IT.PNG", "NL.PNG", "NZ.PNG"]
         gr.Examples(examples=examples, inputs=input_image)
-    gradio_app.launch()
+    gradio_app.launch()