File size: 4,923 Bytes
d758ffc
 
 
 
 
 
 
0246ff9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d758ffc
0246ff9
d758ffc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0246ff9
 
 
 
 
 
 
 
 
d758ffc
 
 
 
 
 
 
0246ff9
d758ffc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0246ff9
 
 
 
 
 
 
d758ffc
 
 
 
 
 
 
 
 
 
 
 
0246ff9
 
 
d758ffc
0246ff9
d758ffc
 
 
0246ff9
 
d758ffc
 
 
 
0246ff9
d758ffc
 
 
 
0246ff9
d758ffc
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
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()