fix: image selection

app.py

@@ -1,37 +1,26 @@
 import gradio as gr
-import numpy as np
 from PIL import Image
 
 from inference import generate_image
 
 
-# Create a square image for the coordinate selector
-def create_selector_image():
-    # Create a white square with black border
-    size = 400
-    border = 2
-    img = np.ones((size, size, 3), dtype=np.uint8) * 255
-    # Add black border
-    img[:border, :] = 0  # top
-    img[-border:, :] = 0  # bottom
-    img[:, :border] = 0  # left
-    img[:, -border:] = 0  # right
-    return Image.fromarray(img)
-
-
 def process_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
     """
     Process the click event on the coordinate selector
     """
-    # Extract coordinates from click event
     x, y = evt.index[0], evt.index[1]
-    # Normalize coordinates to [0, 1]
-    x, y = x / 400, y / 400  # Divide by image size (400x400)
+    x, y = x / 400, y / 400
     print(f"Clicked at coordinates: ({x:.3f}, {y:.3f})")
-    # Generate image using the model
     return generate_image(image_idx, x, y)
 
 
+def update_background(image_idx: int) -> Image.Image:
+    """
+    Update the coordinate selector background based on selected image
+    """
+    return f"imgs/heatmap_{image_idx}.png"
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
         """
@@ -41,7 +30,7 @@ with gr.Blocks() as demo:
     )
 
     with gr.Row():
-        # Left column: Reference images and coordinate selector
+        # Left column: Reference images only
         with gr.Column(scale=1):
             # Radio buttons for image selection
             image_idx = gr.Radio(
@@ -54,18 +43,20 @@ with gr.Blocks() as demo:
             # Display reference images
             gallery = gr.Gallery(
                 value=[
+                    "imgs/pattern_0.png",
                     "imgs/pattern_1.png",
-                    "imgs/pattern_2.png",
                 ],
-                columns=2,
+                columns=1,
                 rows=2,
-                height=300,
+                min_width=600,
                 label="Reference Images",
             )
 
-            # Coordinate selector (now using Image component)
+        # Right column: Coordinate selector and output image
+        with gr.Column(scale=1):
+            # Coordinate selector with dynamic background
             coord_selector = gr.Image(
-                value=create_selector_image(),
+                value="imgs/heatmap_0.png",  # Initial background
                 label="Click to select (x, y) coordinates",
                 show_label=True,
                 interactive=True,
@@ -73,12 +64,11 @@ with gr.Blocks() as demo:
                 width=400,
             )
 
-        # Right column: Generated image
-        with gr.Column(scale=1):
-            output_image = gr.Image(label="Generated Image", height=308, width=328)
+            # Generated image output
+            output_image = gr.Image(label="Generated Image", height=400, width=400)
 
-    # Handle click events
+    # Handle click events and background updates
     coord_selector.select(process_click, inputs=[image_idx], outputs=output_image, trigger_mode="multiple")
+    image_idx.change(update_background, inputs=[image_idx], outputs=[coord_selector])
 
-# Launch the app
 demo.launch()