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 "

Please provide an image.

", "" 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()