app.py

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)
    print(f"Clicked at coordinates: ({x:.3f}, {y:.3f})")
    # Generate image using the model
    return generate_image(image_idx, x, y)


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.
    """
    )

    with gr.Row():
        # Left column: Reference images and coordinate selector
        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",
            )

            # Display reference images
            gallery = gr.Gallery(
                value=[
                    "imgs/pattern_1.png",
                    "imgs/pattern_2.png",
                ],
                columns=2,
                rows=2,
                height=300,
                label="Reference Images",
            )

            # Coordinate selector (now using Image component)
            coord_selector = gr.Image(
                value=create_selector_image(),
                label="Click to select (x, y) coordinates",
                show_label=True,
                interactive=True,
                height=400,
                width=400,
            )

        # Right column: Generated image
        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated Image", height=308, width=328)

    # Handle click events
    coord_selector.select(process_click, inputs=[image_idx], outputs=output_image, trigger_mode="multiple")

# Launch the app
demo.launch()