feat: draw marker before generating image

app.py

@@ -30,17 +30,23 @@ def update_reference_image(choice: int) -> tuple[str, int, str]:
     return image_path, choice, heatmap_path
 
 
-def process_coord_click(image_idx: int, evt: gr.SelectData) -> tuple[Image.Image, Image.Image]:
+def update_marker(image_idx: int, evt: gr.SelectData) -> tuple[Image.Image, tuple[int, int]]:
     """
-    Process the click event on the coordinate selector
-    Returns both the generated image and the marked coordinate selector
+    Update the coordinate selector with the marker
+    Returns the marked image and the coordinates for the next function
     """
     x, y = evt.index[0], evt.index[1]
-    x_norm, y_norm = x / 1155, y / 1155
-    generated_image = generate_image(image_idx, x_norm, y_norm)
     heatmap_path = f"imgs/heatmap_{image_idx}.png"
-    marked_selector = create_marker_overlay(heatmap_path, x, y)
-    return generated_image, marked_selector
+    return create_marker_overlay(heatmap_path, x, y), (x, y)
+
+
+def generate_output_image(image_idx: int, coords: tuple[int, int]) -> Image.Image:
+    """
+    Generate the output image based on the selected coordinates
+    """
+    x, y = coords
+    x_norm, y_norm = x / 1155, y / 1155
+    return generate_image(image_idx, x_norm, y_norm)
 
 
 with gr.Blocks(
@@ -78,14 +84,13 @@ with gr.Blocks(
         # Left column
         with gr.Column(scale=1):
             selected_idx = gr.State(value=0)
+            coords = gr.State()  # Add state for coordinates
 
-            # gr.Markdown("### Select Task")
             with gr.Column(elem_classes="radio-container"):
                 task_select = gr.Radio(
                     choices=["Task 1", "Task 2", "Task 3", "Task 4"],
                     value="Task 1",
                     label="Select Task",
-                    # show_label=False,
                     interactive=True,
                 )
 
@@ -158,11 +163,16 @@ with gr.Blocks(
         outputs=[reference_image, selected_idx, coord_selector],
     )
 
+    # Split the coordinate selection into two events with state passing
     coord_selector.select(
-        process_coord_click,
+        fn=update_marker,
         inputs=[selected_idx],
-        outputs=[output_image, coord_selector],
+        outputs=[coord_selector, coords],
         trigger_mode="multiple",
+    ).then(
+        fn=generate_output_image,
+        inputs=[selected_idx, coords],
+        outputs=output_image,
     )
 
 demo.launch()