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 | |
| # Default Demo Examples | |
| SAMPLE_IMAGES_DIR = "screencoder/data/input" | |
| examples_data = [] | |
| if os.path.exists(SAMPLE_IMAGES_DIR): | |
| sample_files = [f for f in sorted(os.listdir(SAMPLE_IMAGES_DIR)) if f.endswith(('.png', '.jpg', '.jpeg')) and not f.startswith('.')] | |
| for filename in sample_files: | |
| path = os.path.join(SAMPLE_IMAGES_DIR, filename) | |
| prompt = f"Generate a modern UI based on the '{filename}' example, focusing on a clean and intuitive layout." | |
| examples_data.append([path, prompt, path]) | |
| else: | |
| print(f"Warning: Sample images directory not found at {SAMPLE_IMAGES_DIR}. Examples will be empty.") | |
| def process_image_and_prompt(image_np, image_path_from_state, 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>", "" | |
| print(f"With prompt: '{prompt}'") | |
| html_content = generate_html_for_demo(final_image_path, prompt) | |
| 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][2] 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 a Prompt (Optional)") | |
| prompt_input = gr.Textbox( | |
| label="Instructions", | |
| placeholder="e.g., 'Make this a dark theme and change the text.'", | |
| lines=3, | |
| value=examples_data[0][1] if examples_data else "Based on the layout, please fill in appropriate English text and beautify the image blocks." | |
| ) | |
| 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, | |
| inputs=[active_image], | |
| label="Click an example to try it out", | |
| ) | |
| def handle_upload(uploaded_image_np): | |
| """On upload, update image, clear state, and set a generic prompt for user input.""" | |
| default_prompt = "Based on the layout, please fill in appropriate English text and beautify the image blocks." | |
| return uploaded_image_np, None, default_prompt | |
| upload_button.upload( | |
| fn=handle_upload, | |
| inputs=upload_button, | |
| outputs=[active_image, active_image_path_state, prompt_input] | |
| ) | |
| generate_btn.click( | |
| fn=process_image_and_prompt, | |
| inputs=[active_image, active_image_path_state, prompt_input], | |
| outputs=[html_preview, html_code_output], | |
| show_progress="full" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |