feat: optimal latent button

app.py

@@ -4,6 +4,7 @@ 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}
+TASK_OPTIMAL_COORDS = {0: (325, 326), 1: (59, 1126), 2: (47, 102), 3: (497, 933)}
 
 
 def create_marker_overlay(image_path: str, x: int, y: int) -> Image.Image:
@@ -49,6 +50,18 @@ def generate_output_image(image_idx: int, coords: tuple[int, int]) -> Image.Imag
     return generate_image(image_idx, x_norm, y_norm)
 
 
+def find_optimal_latent(image_idx: int) -> tuple[Image.Image, tuple[int, int], Image.Image]:
+    """
+    Simulate clicking at the optimal coordinates for the current task
+    Returns the marked heatmap, coordinates, and generated output image
+    """
+    x, y = TASK_OPTIMAL_COORDS[image_idx]
+    heatmap_path = f"imgs/heatmap_{image_idx}.png"
+    marked_heatmap = create_marker_overlay(heatmap_path, x, y)
+    output_img = generate_output_image(image_idx, (x, y))
+    return marked_heatmap, (x, y), output_img
+
+
 with gr.Blocks(
     css="""
     .radio-container {
@@ -71,6 +84,13 @@ with gr.Blocks(
         background-color: #f8f9fa !important;
         border-radius: 8px !important;
     }
+    .optimal-button {
+        margin-top: 1rem !important;
+        margin-bottom: 1rem !important;
+        width: 200px !important;
+        margin-left: auto !important;
+        margin-right: auto !important;
+    }
 """
 ) as demo:
     gr.Markdown(
@@ -130,6 +150,10 @@ with gr.Blocks(
                     show_fullscreen_button=False,
                 )
 
+            # Add the Find Optimal Latent button with a container for centering
+            with gr.Column():
+                optimal_button = gr.Button("Find Optimal Latent", elem_classes="optimal-button")
+
     # Documentation section
     with gr.Column(elem_classes="documentation"):
         gr.Markdown(
@@ -144,8 +168,7 @@ with gr.Blocks(
         3. **Coordinate Selection**: Click anywhere in the heatmap to specify where in the latent space you want to generate from
         4. **Generation**: The model generates a new image based on your selected coordinates
         
-        ### Sample Results
-        Here are some example outputs from our method:
+        The "Find Optimal Latent" button automatically selects pre-determined optimal coordinates for the current task.
         """
         )
 
@@ -186,4 +209,11 @@ with gr.Blocks(
         outputs=output_image,
     )
 
+    # Add event handler for the optimal latent button
+    optimal_button.click(
+        fn=find_optimal_latent,
+        inputs=[selected_idx],
+        outputs=[coord_selector, coords, output_image],
+    )
+
 demo.launch()