Update app.py

app.py

@@ -29,6 +29,8 @@ device = 'cpu'
 le = LabelEncoder()
 le = joblib.load("SVD/le.gz")
 len_classes = len(le.classes_) + 1
+# Global variable to store predictions for dynamic zoom
+global_predictions = None
 
 class ModelPre(torch.nn.Module):
     def __init__(self):
@@ -59,16 +61,26 @@ cmp = transforms.Compose([
 ])
 
 def predict(input_img):
-    with torch.inference_mode():
-        img = cmp(input_img).unsqueeze(0)
-        res = model(img.to(device))
-        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
+    global global_predictions
+    results, cell_ids = predict(input_img)
+    global_predictions = (results, cell_ids)  # Store predictions globally for zoom functionality
+    return create_map_figure(global_predictions, global_predictions[1])
+
+def zoom_on_prediction(selected_prediction):
+    global global_predictions
+    if global_predictions is None:
+        return None  # No prediction made yet
+    
+    # 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 the updated map with zoom
+    return create_map_figure(global_predictions, global_predictions[1], selected_index=selected_index)
+
+
 
 # Function to get S2 cell polygon
 def get_s2_cell_polygon(cell_id):
@@ -154,12 +166,15 @@ 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", value=None)
         output_map = gr.Plot(label="Predicted Location on Map")
         btn_predict = gr.Button("Predict")
+        selected_prediction = gr.Dropdown(choices=[f"Prediction {i+1}" for i in range(10)], label="Select Prediction to Zoom", value=None)
+
+        # Perform the prediction and plot the initial map
+        btn_predict.click(predict, 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)
+        # Allow the user to zoom in on a selected prediction after the prediction is made
+        selected_prediction.change(zoom_on_prediction, inputs=selected_prediction, outputs=output_map)
 
         examples = ["GB.PNG", "IT.PNG", "NL.PNG", "NZ.PNG"]
         gr.Examples(examples=examples, inputs=input_image)