feat: left images in full

app.py

@@ -1,10 +1,12 @@
+from functools import partial
+
 import gradio as gr
 from PIL import Image
 
 from inference import generate_image
 
 
-def process_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
+def process_coord_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
     """
     Process the click event on the coordinate selector
     """
@@ -14,43 +16,46 @@ def process_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
     return generate_image(image_idx, x, y)
 
 
-def update_background(image_idx: int) -> Image.Image:
+def process_image_select(evt: gr.SelectData, idx: int) -> tuple[int, str]:
     """
-    Update the coordinate selector background based on selected image
+    Process the reference image selection
+    Returns the selected image index and corresponding heatmap
     """
-    return f"imgs/heatmap_{image_idx}.png"
+    return idx, f"imgs/heatmap_{idx}.png"
 
 
 with gr.Blocks() as demo:
     gr.Markdown(
         """
     # Interactive Image Generation
-    Choose a reference image and click on the coordinate selector to generate a new image.
+    Click on a reference image to select it, then click on the coordinate selector to generate a new image.
     """
     )
 
     with gr.Row():
-        # Left column: Reference images only
+        # Left column: Interactive reference images
         with gr.Column(scale=1):
-            # Radio buttons for image selection
-            image_idx = gr.Radio(
-                choices=list(range(2)),
-                value=0,
-                label="Select Reference Image",
-                type="index",
-            )
+            # State variable to track selected image index
+            selected_idx = gr.State(value=0)
 
-            # Display reference images
-            gallery = gr.Gallery(
-                value=[
-                    "imgs/pattern_0.png",
-                    "imgs/pattern_1.png",
-                ],
-                columns=1,
-                rows=2,
-                min_width=600,
-                label="Reference Images",
-            )
+            # Two separate Image components for reference images
+            with gr.Column():
+                image_0 = gr.Image(
+                    value="imgs/pattern_0.png",
+                    label="Task 1",
+                    show_label=False,
+                    interactive=True,
+                    height=300,
+                    width=450,
+                )
+                image_1 = gr.Image(
+                    value="imgs/pattern_1.png",
+                    label="Task 2",
+                    show_label=False,
+                    interactive=True,
+                    height=300,
+                    width=450,
+                )
 
         # Right column: Coordinate selector and output image
         with gr.Column(scale=1):
@@ -65,10 +70,15 @@ with gr.Blocks() as demo:
             )
 
             # Generated image output
-            output_image = gr.Image(label="Generated Image", height=400, width=400)
+            output_image = gr.Image(label="Generated Output", height=400, width=400)
 
-    # 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])
+    # Handle image selection for each reference image
+    image_0.select(partial(process_image_select, idx=0), outputs=[selected_idx, coord_selector])
+    image_1.select(partial(process_image_select, idx=1), outputs=[selected_idx, coord_selector])
+
+    # Handle coordinate selection
+    coord_selector.select(
+        process_coord_click, inputs=[selected_idx], outputs=output_image, trigger_mode="multiple"
+    )
 
 demo.launch()