Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import tempfile | |
| import cv2 | |
| import numpy as np | |
| from screencoder.main import generate_html_for_demo | |
| # Manually defined examples | |
| examples_data = [ | |
| [ | |
| "screencoder/data/input/test1.png", | |
| "", | |
| "", | |
| "", | |
| "", | |
| "screencoder/data/input/test1.png" | |
| ], | |
| [ | |
| "screencoder/data/input/test2.png", | |
| "", | |
| "", | |
| "", | |
| "", | |
| "screencoder/data/input/test2.png" | |
| ], | |
| [ | |
| "screencoder/data/input/test3.png", | |
| "", | |
| "", | |
| "", | |
| "", | |
| "screencoder/data/input/test3.png" | |
| ], | |
| ] | |
| def process_image_and_prompt(image_np, image_path_from_state, sidebar_prompt, header_prompt, navigation_prompt, main_content_prompt): | |
| final_image_path = "" | |
| is_temp_file = False | |
| if image_path_from_state: | |
| final_image_path = image_path_from_state | |
| print(f"Processing example image from: {final_image_path}") | |
| elif image_np is not None: | |
| is_temp_file = True | |
| with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: | |
| image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR) | |
| cv2.imwrite(tmp.name, image_bgr) | |
| final_image_path = tmp.name | |
| print(f"Processing uploaded image from temporary file: {final_image_path}") | |
| else: | |
| return "<html><body><h1 style='font-family: sans-serif; text-align: center; margin-top: 40px;'>Please provide an image.</h1></body></html>", "" | |
| instructions = { | |
| "sidebar": sidebar_prompt, | |
| "header": header_prompt, | |
| "navigation": navigation_prompt, | |
| "main content": main_content_prompt | |
| } | |
| print(f"With instructions: {instructions}") | |
| html_content = generate_html_for_demo(final_image_path, instructions) | |
| if is_temp_file: | |
| os.unlink(final_image_path) | |
| return html_content, html_content | |
| with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), fill_height=True) as demo: | |
| active_image_path_state = gr.State(value=examples_data[0][5] if examples_data else None) | |
| gr.Markdown("# ScreenCoder: Screenshot to Code") | |
| with gr.Row(equal_height=True): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Step 1: Provide an Image") | |
| active_image = gr.Image( | |
| type="numpy", | |
| height=300, | |
| value=examples_data[0][0] if examples_data else None | |
| ) | |
| upload_button = gr.UploadButton("Click to Upload or Drag-and-Drop", file_types=["image"], variant="primary") | |
| gr.Markdown("### Step 2: Write Prompts (Optional)") | |
| with gr.Accordion("Component-specific Prompts", open=True): | |
| sidebar_prompt = gr.Textbox(label="Sidebar Prompt", placeholder="Instructions for the sidebar...", value="") | |
| header_prompt = gr.Textbox(label="Header Prompt", placeholder="Instructions for the header...", value="") | |
| navigation_prompt = gr.Textbox(label="Navigation Prompt", placeholder="Instructions for the navigation...", value="") | |
| main_content_prompt = gr.Textbox(label="Main Content Prompt", placeholder="Instructions for the main content...", value="") | |
| generate_btn = gr.Button("Generate HTML", variant="primary", scale=2) | |
| with gr.Column(scale=2): | |
| with gr.Tabs(): | |
| with gr.TabItem("Preview"): | |
| html_preview = gr.HTML(label="Live Preview", elem_id="html-preview") | |
| with gr.TabItem("Code"): | |
| html_code_output = gr.Code(label="Generated Code", language="html", elem_id="html-code") | |
| if examples_data: | |
| gr.Examples( | |
| examples=examples_data, | |
| fn=lambda *args: args, # Simply return all inputs | |
| inputs=[active_image, sidebar_prompt, header_prompt, navigation_prompt, main_content_prompt, active_image_path_state], | |
| outputs=[active_image, sidebar_prompt, header_prompt, navigation_prompt, main_content_prompt, active_image_path_state], | |
| label="Click an example to try it out", | |
| cache_examples=False, | |
| ) | |
| def handle_upload(uploaded_image_np): | |
| """On upload, update image, clear state, and set empty prompts.""" | |
| return uploaded_image_np, None, "", "", "", "" | |
| upload_button.upload( | |
| fn=handle_upload, | |
| inputs=upload_button, | |
| outputs=[active_image, active_image_path_state, sidebar_prompt, header_prompt, navigation_prompt, main_content_prompt] | |
| ) | |
| generate_btn.click( | |
| fn=process_image_and_prompt, | |
| inputs=[active_image, active_image_path_state, sidebar_prompt, header_prompt, navigation_prompt, main_content_prompt], | |
| outputs=[html_preview, html_code_output], | |
| show_progress="full" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |