Update app.py

app.py

@@ -80,17 +80,21 @@ def get_s2_cell_polygon(cell_id):
     vertices.append(vertices[0])  # Close the polygon
     return vertices
 
-def create_map_figure(predictions, cell_ids):
+def create_map_figure(predictions, cell_ids, selected_index=None):
     fig = go.Figure()
 
     # Assign colors based on rank
     colors = ['rgba(0, 255, 0, 0.2)'] * 3 + ['rgba(255, 255, 0, 0.2)'] * 7
+    zoom_level = 1
+    center_lat = None
+    center_lon = None
 
     for rank, cell_id in enumerate(cell_ids):
         cell_id = int(float(cell_id))
         polygon = get_s2_cell_polygon(cell_id)
         lats, lons = zip(*polygon)
         color = colors[rank]
+        
         fig.add_trace(go.Scattermapbox(
             lat=lats,
             lon=lons,
@@ -101,17 +105,23 @@ def create_map_figure(predictions, cell_ids):
             name=f'Prediction {rank + 1}',  # Updated label
         ))
 
+        # Set zoom based on the selected index
+        if selected_index is not None and rank == selected_index:
+            zoom_level = 10  # Adjust zoom level
+            center_lat = np.mean(lats)
+            center_lon = np.mean(lons)
+
     fig.update_layout(
-        mapbox_style="open-street-map",  # Change this line to use 'light' style
+        mapbox_style="open-street-map",
         hovermode='closest',
         mapbox=dict(
             bearing=0,
             center=go.layout.mapbox.Center(
-                lat=np.mean(lats),
-                lon=np.mean(lons)
+                lat=center_lat if center_lat else np.mean(lats),
+                lon=center_lon if center_lon else np.mean(lons)
             ),
             pitch=0,
-            zoom=1
+            zoom=zoom_level  # Zoom in if an index is selected
         ),
     )
 
@@ -124,20 +134,24 @@ def create_label_output(predictions):
     fig = create_map_figure(results, cell_ids)
     return fig
 
-# Predict and plot function
-def predict_and_plot(input_img):
+# Update the predict_and_plot function to handle zoom on selection
+def predict_and_plot(input_img, selected_prediction):
     predictions = predict(input_img)
-    return create_label_output(predictions)
+    return create_map_figure(predictions, predictions[1], selected_index=selected_prediction)
+
 
 
 # Gradio app definition
 with gr.Blocks() as gradio_app:
     with gr.Column():
         input_image = gr.Image(label="Upload an Image", type="pil")
+        selected_prediction = gr.Dropdown(choices=[f"Prediction {i+1}" for i in range(10)], label="Select Prediction to Zoom")
         output_map = gr.Plot(label="Predicted Location on Map")
         btn_predict = gr.Button("Predict")
 
-        btn_predict.click(predict_and_plot, inputs=input_image, outputs=output_map)
+        # Update click function to include selected prediction
+        btn_predict.click(predict_and_plot, inputs=[input_image, selected_prediction], 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()