Spaces:
Running
Running
| import gradio as gr | |
| 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: | |
| """Creates an image with a marker at the specified coordinates""" | |
| base_image = Image.open(image_path) | |
| marked_image = base_image.copy() | |
| draw = ImageDraw.Draw(marked_image) | |
| marker_size = 10 | |
| marker_color = "red" | |
| draw.line([x - marker_size, y, x + marker_size, y], fill=marker_color, width=2) | |
| draw.line([x, y - marker_size, x, y + marker_size], fill=marker_color, width=2) | |
| return marked_image | |
| def update_reference_image(choice: int) -> tuple[str, int, str]: | |
| image_path = f"imgs/pattern_{choice}.png" | |
| heatmap_path = f"imgs/heatmap_{choice}.png" | |
| return image_path, choice, heatmap_path | |
| def update_marker(image_idx: int, evt: gr.SelectData) -> tuple[Image.Image, tuple[int, int]]: | |
| x, y = evt.index[0], evt.index[1] | |
| heatmap_path = f"imgs/heatmap_{image_idx}.png" | |
| return create_marker_overlay(heatmap_path, x, y), (x, y) | |
| def generate_output_image(image_idx: int, coords: tuple[int, int]) -> Image.Image: | |
| x, y = coords | |
| x_norm, y_norm = x / 1155, y / 1155 | |
| return generate_image(image_idx, x_norm, y_norm) | |
| def find_optimal_latent(image_idx: int) -> tuple[Image.Image, tuple[int, int], Image.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=""" | |
| .container { | |
| max-width: 1000px !important; | |
| width: 100% !important; | |
| margin-left: auto !important; | |
| margin-right: auto !important; | |
| padding: 0 1rem !important; | |
| } | |
| .diagram-container { | |
| width: 80% !important; | |
| max-width: 800px !important; | |
| margin-left: auto !important; | |
| margin-right: auto !important; | |
| } | |
| .diagram-container img { | |
| width: 100% !important; | |
| height: auto !important; | |
| object-fit: contain !important; | |
| } | |
| .radio-container { | |
| width: 100% !important; | |
| max-width: 450px !important; | |
| margin-left: auto !important; | |
| margin-right: auto !important; | |
| } | |
| .image-container { | |
| width: 100% !important; | |
| aspect-ratio: 1 !important; | |
| max-width: 450px !important; | |
| margin-left: auto !important; | |
| margin-right: auto !important; | |
| } | |
| .image-container img { | |
| width: 100% !important; | |
| height: 100% !important; | |
| object-fit: contain !important; | |
| } | |
| .coordinate-container { | |
| width: 100% !important; | |
| aspect-ratio: 1 !important; | |
| position: relative !important; | |
| max-width: 600px !important; | |
| margin-left: auto !important; | |
| margin-right: auto !important; | |
| } | |
| .coordinate-container img { | |
| width: 100% !important; | |
| height: 100% !important; | |
| object-fit: contain !important; | |
| } | |
| .button-container { | |
| display: flex !important; | |
| justify-content: center !important; | |
| width: 100% !important; | |
| max-width: 600px !important; | |
| margin-left: auto !important; | |
| margin-right: auto !important; | |
| } | |
| .documentation { | |
| margin-top: 2rem !important; | |
| padding: 1rem !important; | |
| background-color: #f8f9fa !important; | |
| border-radius: 8px !important; | |
| } | |
| .optimal-button { | |
| margin-top: 1rem !important; | |
| margin-bottom: 1rem !important; | |
| width: 200px !important; | |
| } | |
| """ | |
| ) as demo: | |
| with gr.Column(elem_classes="container"): | |
| gr.Markdown( | |
| """ | |
| # Interactive Image Generation | |
| ## Method Overview | |
| This interactive demo showcases our novel image generation method that uses coordinate-based control. | |
| The process allows precise control over generated patterns through a coordinate-conditioning mechanism. | |
| """ | |
| ) | |
| with gr.Column(elem_classes="diagram-container"): | |
| gr.Image( | |
| value="imgs/lpn_diagram.png", | |
| show_label=False, | |
| interactive=False, | |
| show_download_button=False, | |
| show_fullscreen_button=False, | |
| ) | |
| gr.Markdown( | |
| """ | |
| ### How to Use | |
| 1. Choose a pattern generation task using the radio buttons | |
| 2. View the target pattern for your selected task | |
| 3. Click anywhere in the heatmap to specify coordinates in the latent space | |
| 4. See the generated image based on your selection | |
| Use the "Find Optimal Latent" button to automatically select pre-determined optimal coordinates. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| selected_idx = gr.State(value=0) | |
| coords = gr.State() | |
| with gr.Column(elem_classes="radio-container"): | |
| task_select = gr.Radio( | |
| choices=["Task 1", "Task 2", "Task 3", "Task 4"], | |
| value="Task 1", | |
| label="Select Task", | |
| interactive=True, | |
| ) | |
| gr.Markdown("### Reference Pattern") | |
| with gr.Column(elem_classes="image-container"): | |
| reference_image = gr.Image( | |
| value="imgs/pattern_0.png", | |
| show_label=False, | |
| interactive=False, | |
| show_download_button=False, | |
| show_fullscreen_button=False, | |
| ) | |
| gr.Markdown("### Generated Output") | |
| with gr.Column(elem_classes="image-container"): | |
| output_image = gr.Image( | |
| show_label=False, | |
| interactive=False, | |
| show_download_button=False, | |
| show_fullscreen_button=False, | |
| ) | |
| with gr.Column(scale=4): | |
| gr.Markdown("### Coordinate Selector") | |
| gr.Markdown( | |
| "Click anywhere in the image below to select (x, y) coordinates in the latent space" | |
| ) | |
| with gr.Column(elem_classes="coordinate-container"): | |
| coord_selector = gr.Image( | |
| value="imgs/heatmap_0.png", | |
| show_label=False, | |
| interactive=False, | |
| sources=[], | |
| container=True, | |
| show_download_button=False, | |
| show_fullscreen_button=False, | |
| ) | |
| with gr.Column(elem_classes="button-container"): | |
| optimal_button = gr.Button("Find Optimal Latent", elem_classes="optimal-button") | |
| with gr.Column(elem_classes="documentation"): | |
| gr.Markdown( | |
| """ | |
| ### Technical Details | |
| Our approach uses a novel coordinate-conditioning mechanism that allows precise control over the generated patterns. | |
| The heatmap visualization shows the distribution of pattern characteristics across the latent space. | |
| For more information, please refer to our [paper](https://arxiv.org/pdf/2411.08706) or GitHub [repository](https://github.com/clement-bonnet/lpn). | |
| """ | |
| ) | |
| # Event handlers | |
| task_select.change( | |
| fn=lambda x: update_reference_image(TASK_TO_INDEX[x]), | |
| inputs=[task_select], | |
| outputs=[reference_image, selected_idx, coord_selector], | |
| ) | |
| coord_selector.select( | |
| fn=update_marker, | |
| inputs=[selected_idx], | |
| outputs=[coord_selector, coords], | |
| trigger_mode="multiple", | |
| ).then( | |
| fn=generate_output_image, | |
| inputs=[selected_idx, coords], | |
| outputs=output_image, | |
| ) | |
| optimal_button.click( | |
| fn=find_optimal_latent, | |
| inputs=[selected_idx], | |
| outputs=[coord_selector, coords, output_image], | |
| ) | |
| demo.launch() | |