feat: add marker on the heatmap

app.py

@@ -1,12 +1,33 @@
 import gradio as gr
-from PIL import Image
+from PIL import Image, ImageDraw
 
 from inference import generate_image
 
-
 TASK_TO_INDEX = {"Task 1": 0, "Task 2": 1, "Task 3": 2, "Task 4": 3}
 
 
+def create_marker_overlay(image_path: str, x: int, y: int) -> Image.Image:
+    """
+    Creates an image with a marker at the specified coordinates
+    """
+    # Load the base image
+    base_image = Image.open(image_path)
+
+    # Create a copy to draw on
+    marked_image = base_image.copy()
+    draw = ImageDraw.Draw(marked_image)
+
+    # Define marker properties
+    marker_size = 10
+    marker_color = "red"
+
+    # Draw marker
+    draw.line([x - marker_size, y, x + marker_size, y], fill=marker_color, width=2)
+    draw.line([x, y - marker_size, x, y + marker_size], fill=marker_color, width=2)
+
+    return marked_image
+
+
 def update_reference_image(choice: int) -> tuple[str, int, str]:
     """
     Update the reference image display based on radio button selection
@@ -17,13 +38,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) -> Image.Image:
+def process_coord_click(image_idx: int, evt: gr.SelectData) -> tuple[Image.Image, Image.Image]:
     """
     Process the click event on the coordinate selector
+    Returns both the generated image and the marked coordinate selector
     """
     x, y = evt.index[0], evt.index[1]
-    x, y = x / 1155, y / 1155  # Normalize the coordinates
-    return generate_image(image_idx, x, y)
+    # Create normalized coordinates for generation
+    x_norm, y_norm = x / 1155, y / 1155
+
+    # Generate the output image
+    generated_image = generate_image(image_idx, x_norm, y_norm)
+
+    # Create marked coordinate selector
+    heatmap_path = f"imgs/heatmap_{image_idx}.png"
+    marked_selector = create_marker_overlay(heatmap_path, x, y)
+
+    return generated_image, marked_selector
 
 
 with gr.Blocks(
@@ -109,9 +140,12 @@ with gr.Blocks(
         outputs=[reference_image, selected_idx, coord_selector],
     )
 
-    # Handle coordinate selection
+    # Handle coordinate selection - now updates both output image and coord_selector
     coord_selector.select(
-        process_coord_click, inputs=[selected_idx], outputs=output_image, trigger_mode="multiple"
+        process_coord_click,
+        inputs=[selected_idx],
+        outputs=[output_image, coord_selector],
+        trigger_mode="multiple",
     )
 
 demo.launch()