feat: radio button

app.py

@@ -1,11 +1,19 @@
-from functools import partial
-
 import gradio as gr
 from PIL import Image
 
 from inference import generate_image
 
 
+def update_reference_image(choice: int) -> tuple[str, int, str]:
+    """
+    Update the reference image display based on radio button selection
+    Returns the image path, selected index, and corresponding heatmap
+    """
+    image_path = f"imgs/pattern_{choice}.png"
+    heatmap_path = f"imgs/heatmap_{choice}.png"
+    return image_path, choice, heatmap_path
+
+
 def process_coord_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
     """
     Process the click event on the coordinate selector
@@ -16,53 +24,43 @@ def process_coord_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
     return generate_image(image_idx, x, y)
 
 
-def process_image_select(evt: gr.SelectData, idx: int) -> tuple[int, str]:
-    """
-    Process the reference image selection
-    Returns the selected image index and corresponding heatmap
-    """
-    return idx, f"imgs/heatmap_{idx}.png"
-
-
 with gr.Blocks() as demo:
     gr.Markdown(
         """
     # Interactive Image Generation
-    Click on a reference image to select it, then click on the coordinate selector to generate a new image.
+    Select a task using the radio buttons, then click on the coordinate selector to generate a new image.
     """
     )
 
     with gr.Row():
-        # Left column: Interactive reference images
+        # Left column: Radio selection and reference image
         with gr.Column(scale=1):
             # State variable to track selected image index
             selected_idx = gr.State(value=0)
 
-            # 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,
-                )
+            # Radio buttons for task selection
+            task_select = gr.Radio(
+                choices=["Task 1", "Task 2"],
+                value="Task 1",
+                label="Select Task",
+                interactive=True,
+            )
+
+            # Single reference image component that updates based on selection
+            reference_image = gr.Image(
+                value="imgs/pattern_0.png",
+                show_label=False,
+                interactive=False,  # No need for click interaction now
+                height=300,
+                width=450,
+            )
 
         # Right column: Coordinate selector and output image
         with gr.Column(scale=1):
             # Coordinate selector with dynamic background
             coord_selector = gr.Image(
-                value="imgs/heatmap_0.png",  # Initial background
-                label="Click to select (x, y) coordinates",
+                value="imgs/heatmap_0.png",
+                label="Click to select (x, y) coordinates in the latent space",
                 show_label=True,
                 interactive=True,
                 height=400,
@@ -72,9 +70,12 @@ with gr.Blocks() as demo:
             # Generated image output
             output_image = gr.Image(label="Generated Output", height=400, width=400)
 
-    # 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 radio button selection
+    task_select.change(
+        fn=lambda x: update_reference_image(0 if x == "Task 1" else 1),
+        inputs=[task_select],
+        outputs=[reference_image, selected_idx, coord_selector],
+    )
 
     # Handle coordinate selection
     coord_selector.select(