Update app.py

app.py

@@ -85,7 +85,7 @@ def create_map_figure(predictions, cell_ids, selected_index=None):
 
     # Assign colors based on rank
     colors = ['rgba(0, 255, 0, 0.2)'] * 3 + ['rgba(255, 255, 0, 0.2)'] * 7
-    zoom_level = 1
+    zoom_level = 1  # Default zoom level
     center_lat = None
     center_lon = None
 
@@ -94,7 +94,8 @@ def create_map_figure(predictions, cell_ids, selected_index=None):
         polygon = get_s2_cell_polygon(cell_id)
         lats, lons = zip(*polygon)
         color = colors[rank]
-        
+
+        # Draw S2 cell polygon
         fig.add_trace(go.Scattermapbox(
             lat=lats,
             lon=lons,
@@ -102,15 +103,16 @@ def create_map_figure(predictions, cell_ids, selected_index=None):
             fill='toself',
             fillcolor=color,
             line=dict(color='blue'),
-            name=f'Prediction {rank + 1}',  # Updated label
+            name=f'Prediction {rank + 1}',
         ))
 
-        # Set zoom based on the selected index
+        # Adjust zoom level if selected prediction is found
         if selected_index is not None and rank == selected_index:
-            zoom_level = 10  # Adjust zoom level
+            zoom_level = 10  # Adjust the zoom level to your liking
             center_lat = np.mean(lats)
             center_lon = np.mean(lons)
 
+    # Update map layout
     fig.update_layout(
         mapbox_style="open-street-map",
         hovermode='closest',
@@ -121,7 +123,7 @@ def create_map_figure(predictions, cell_ids, selected_index=None):
                 lon=center_lon if center_lon else np.mean(lons)
             ),
             pitch=0,
-            zoom=zoom_level  # Zoom in if an index is selected
+            zoom=zoom_level  # Zoom in based on selection
         ),
     )
 
@@ -134,10 +136,17 @@ def create_label_output(predictions):
     fig = create_map_figure(results, cell_ids)
     return fig
 
-# 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_map_figure(predictions, predictions[1], selected_index=selected_prediction)
+
+    # Convert dropdown selection into an index (Prediction 1 corresponds to index 0, etc.)
+    if selected_prediction is not None:
+        selected_index = int(selected_prediction.split()[-1]) - 1  # Extract index from "Prediction X"
+    else:
+        selected_index = None  # No selection, default view
+
+    return create_map_figure(predictions, predictions[1], selected_index=selected_index)
+
 
 
 
@@ -145,7 +154,7 @@ def predict_and_plot(input_img, selected_prediction):
 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")
+        selected_prediction = gr.Dropdown(choices=[f"Prediction {i+1}" for i in range(10)], label="Select Prediction to Zoom", value=None)
         output_map = gr.Plot(label="Predicted Location on Map")
         btn_predict = gr.Button("Predict")