fix: x y corrdinates

app.py

@@ -22,18 +22,26 @@ def process_coord_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
     Process the click event on the coordinate selector
     """
     x, y = evt.index[0], evt.index[1]
-    x, y = x / 400, y / 400
-    print(f"Clicked at coordinates: ({x:.3f}, {y:.3f})")
+    x, y = x / 1155, y / 1155  # Normalize the coordinates
     return generate_image(image_idx, x, y)
 
 
 with gr.Blocks(
     css="""
-    .radio-container { 
+    .radio-container {
         width: 450px !important;
         margin-left: auto !important;
         margin-right: auto !important;
     }
+    .coordinate-container {
+        width: 600px !important;
+        height: 600px !important;
+    }
+    .coordinate-container img {
+        width: 100% !important;
+        height: 100% !important;
+        object-fit: contain !important;
+    }
 """
 ) as demo:
     gr.Markdown(
@@ -44,7 +52,7 @@ with gr.Blocks(
     )
 
     with gr.Row():
-        # Left column: Radio selection and reference image
+        # Left column: Radio selection, reference image, and output
         with gr.Column(scale=1):
             # State variable to track selected image index
             selected_idx = gr.State(value=0)
@@ -58,7 +66,7 @@ with gr.Blocks(
                     interactive=True,
                 )
 
-            # Single reference image component that updates based on selection
+            # Reference image component that updates based on selection
             reference_image = gr.Image(
                 value="imgs/pattern_0.png",
                 show_label=False,
@@ -67,20 +75,20 @@ with gr.Blocks(
                 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",
-                label="Click to select (x, y) coordinates in the latent space",
-                show_label=True,
-                interactive=True,
-                height=400,
-                width=400,
-            )
+            # Generated image output moved below reference image
+            output_image = gr.Image(label="Generated Output", height=300, width=450)
 
-            # Generated image output
-            output_image = gr.Image(label="Generated Output", height=400, width=400)
+        # Right column: Larger coordinate selector
+        with gr.Column(scale=1):
+            # Coordinate selector with container class for proper scaling
+            with gr.Column(elem_classes="coordinate-container"):
+                coord_selector = gr.Image(
+                    value="imgs/heatmap_0.png",
+                    label="Click to select (x, y) coordinates in the latent space",
+                    show_label=True,
+                    interactive=True,
+                    container=True,
+                )
 
     # Handle radio button selection
     task_select.change(

inference.py

@@ -82,6 +82,8 @@ generate_output_from_context = jax.jit(
 def generate_image(image_idx: int, x: float, y: float, eps: float = 1e-4) -> Image.Image:
     # Create the input image
     input = jnp.zeros(16, int).at[BLUE_LOCATION_INPUTS[image_idx]].set(1).reshape(4, 4)
+    # Inverse the y coordinate
+    y = 1 - y
     # Ensure x and y are in [eps, 1 - eps]
     x = min(1 - eps, max(eps, x))
     y = min(1 - eps, max(eps, y))