Spaces:

Jimmyzheng-10
/

ScreenCoder

Running

App Files Files Community

Jimmyzheng-10 commited on 17 days ago

Commit

1cc14d1

1 Parent(s): 0246ff9

update

Browse files

Files changed (7) hide show

app.py +217 -57
screencoder/block_parsor.py +0 -5
screencoder/html_generator.py +114 -21
screencoder/image_box_detection.py +1 -5
screencoder/image_replacer.py +11 -16
screencoder/main.py +5 -3
screencoder/mapping.py +0 -2

app.py CHANGED Viewed

@@ -3,9 +3,15 @@ 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",
@@ -33,97 +39,251 @@ examples_data = [
     ],
 ]
-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()

 import tempfile
 import cv2
 import numpy as np
+import urllib.parse
 from screencoder.main import generate_html_for_demo
+from PIL import Image
+import shutil
+import html
+import base64
+from bs4 import BeautifulSoup
+# Predefined examples
 examples_data = [
     [
         "screencoder/data/input/test1.png",
     ],
 ]
+TAILWIND_SCRIPT = "<script src='https://cdn.jsdelivr.net/npm/@tailwindcss/browser@4'></script>"
+def image_to_data_url(image_path):
+    """Convert an image file to a data URL for embedding in HTML."""
+    try:
+        with open(image_path, 'rb') as img_file:
+            img_data = img_file.read()
+            # Detect image type from file extension
+            ext = os.path.splitext(image_path)[1].lower()
+            mime_type = {
+                '.png': 'image/png',
+                '.jpg': 'image/jpeg',
+                '.jpeg': 'image/jpeg',
+                '.gif': 'image/gif',
+                '.webp': 'image/webp'
+            }.get(ext, 'image/png')
+            encoded = base64.b64encode(img_data).decode('utf-8')
+            return f'data:{mime_type};base64,{encoded}'
+    except Exception as e:
+        print(f"Error converting image to data URL: {e}")
+        return None
+def render_preview(code: str, width: int, height: int, scale: float) -> str:
+    """
+    Preview renderer with both width and height control for the inner canvas.
+    """
+    safe_code = html.escape(code).replace("'", "&apos;")
+    iframe_html = f"""
+    <div style="width: 100%; max-width: 1920px; margin: 0 auto; overflow-x: auto; overflow-y: hidden;">
+        <div style="
+            width: 1920px;
+            height: 1000px;
+            margin: 0 auto;
+            display: flex;
+            justify-content: center;
+            align-items: center;
+            border: 1px solid #ddd;
+            overflow: hidden;
+            background: #f9fafb;
+            position: relative;
+            box-shadow: 0 4px 12px rgba(0,0,0,0.1);">
+            <div style="
+                width: {width}px;
+                height: {height}px;
+                transform: scale({scale});
+                transform-origin: top left;
+                border: none;
+                position: relative;">
+                <iframe
+                    style="width: 100%; height: 100%; border: none; display: block;"
+                    srcdoc='{safe_code}'>
+                </iframe>
+            </div>
+        </div>
+    </div>
+    """
+    return iframe_html
+def process_and_generate(image_np, image_path_from_state, sidebar_prompt, header_prompt, navigation_prompt, main_content_prompt):
+    """
+    Main processing pipeline: takes an image, generates code, creates a downloadable
+    package, and returns the initial preview and code outputs.
+    """
     final_image_path = ""
     is_temp_file = False
     if image_path_from_state:
         final_image_path = image_path_from_state
     elif image_np is not None:
         is_temp_file = True
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+            cv2.imwrite(tmp.name, cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR))
             final_image_path = tmp.name
     else:
+        return "No image provided.", "Please upload or select an image.", gr.update(visible=False), None
     instructions = {
+        "sidebar": sidebar_prompt, "header": header_prompt,
+        "navigation": navigation_prompt, "main content": main_content_prompt
     }
+    html_content, run_id = generate_html_for_demo(final_image_path, instructions)
+    if not run_id: # Handle potential errors from the generator
+        return "Generation failed.", f"Error: {html_content}", gr.update(visible=False), None
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    soup = BeautifulSoup(html_content, 'html.parser')
+    print(f"Processing HTML for run_id: {run_id}")
+    for img in soup.find_all('img'):
+        if img.get('src') and not img['src'].startswith(('http', 'data:')):
+            original_src = img['src']
+            print(f"Processing image: {original_src}")
+            if original_src.startswith('/'):
+                img_path = os.path.join(base_dir, original_src.lstrip('/'))
+            else:
+                possible_paths = [
+                    os.path.join(base_dir, 'screencoder', 'data', 'output', run_id, original_src),
+                    os.path.join(base_dir, 'screencoder', 'data', 'output', run_id, 'cropped_images', original_src),
+                    os.path.join(base_dir, 'screencoder', 'data', 'output', run_id, f'cropped_images_{run_id}', original_src),
+                    os.path.join(base_dir, 'screencoder', 'data', 'output', run_id, os.path.basename(original_src)),
+                    os.path.join(base_dir, 'screencoder', 'data', 'output', run_id, original_src.lstrip('./').lstrip('../')),
+                ]
+                img_path = None
+                for path in possible_paths:
+                    if os.path.exists(path):
+                        img_path = path
+                        print(f"Found image at: {path}")
+                        break
+                if not img_path:
+                    print(f"Could not find image: {original_src}")
+                    # Try to convert to data URL as last resort
+                    fallback_path = os.path.join(base_dir, 'screencoder', 'data', 'output', run_id, original_src)
+                    data_url = image_to_data_url(fallback_path)
+                    if data_url:
+                        print(f"Converted to data URL: {original_src}")
+                        img['src'] = data_url
+                    else:
+                        img['src'] = original_src
+                    continue
+            # Check if the image file actually exists
+            if img_path and os.path.exists(img_path):
+                print(f"Found image: {img_path}")
+                # Convert to base64 data URL for better iframe compatibility
+                data_url = image_to_data_url(img_path)
+                if data_url:
+                    print(f"Converted to data URL: {original_src}")
+                    img['src'] = data_url
+                else:
+                    # Fallback to Gradio file path
+                    img['src'] = f'/file={img_path}'
+            else:
+                print(f"Image not found: {img_path}")
+                data_url = image_to_data_url(img_path)
+                if data_url:
+                    print(f"Converted to data URL: {original_src}")
+                    img['src'] = data_url
+                else:
+                    img['src'] = original_src
+    html_content = str(soup)
+    output_dir = os.path.join(base_dir, 'screencoder', 'data', 'output', run_id)
+    packages_dir = os.path.join(base_dir, 'screencoder', 'data', 'packages')
+    os.makedirs(packages_dir, exist_ok=True)
+    shutil.make_archive(os.path.join(packages_dir, run_id), 'zip', output_dir)
+    package_path = os.path.join(packages_dir, f'{run_id}.zip')
+    package_url = f'/file={package_path}'
     if is_temp_file:
         os.unlink(final_image_path)
+    initial_preview = render_preview(html_content, 1280, 600, 0.7)
+    return initial_preview, html_content, gr.update(value=package_url, visible=True)
+with gr.Blocks(head=TAILWIND_SCRIPT, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky")) as demo:
     gr.Markdown("# ScreenCoder: Screenshot to Code")
+    with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### Step 1: Provide an Image")
+            active_image = gr.Image(type="numpy", height=400, value=examples_data[0][0])
+            upload_button = gr.UploadButton("Click to Upload", file_types=["image"], variant="primary")
             gr.Markdown("### Step 2: Write Prompts (Optional)")
+            with gr.Accordion("Component-specific Prompts", open=False):
+                sidebar_prompt = gr.Textbox(label="Sidebar", placeholder="Instructions for the sidebar...")
+                header_prompt = gr.Textbox(label="Header", placeholder="Instructions for the header...")
+                navigation_prompt = gr.Textbox(label="Navigation", placeholder="Instructions for the navigation...")
+                main_content_prompt = gr.Textbox(label="Main Content", placeholder="Instructions for the main content...")
+            generate_btn = gr.Button("Generate HTML", variant="primary")
         with gr.Column(scale=2):
             with gr.Tabs():
                 with gr.TabItem("Preview"):
+                    with gr.Row():
+                        scale_slider = gr.Slider(0.2, 1.5, value=0.7, step=0.05, label="Zoom")
+                        width_slider = gr.Slider(400, 1920, value=1280, step=100, label="Canvas Width (px)")
+                        height_slider = gr.Slider(300, 1080, value=600, step=50, label="Canvas Height (px)")
+                    html_preview = gr.HTML(label="Rendered HTML", show_label=False, container=False)
                 with gr.TabItem("Code"):
+                    html_code_output = gr.Code(label="Generated HTML", language="html")
+            download_button = gr.Button("⬇️ Download Package", visible=False, variant="secondary")
+    gr.Examples(
+        examples=examples_data,
+        fn=lambda *args: args[0],
+        inputs=[gr.State(examples_data[0][0])],
+        outputs=[active_image],
+        cache_examples=False,
+    )
+    active_image_path_state = gr.State(examples_data[0][5])
+    def handle_example_click(img_path):
+        return img_path, img_path
+    demo.load(
+        lambda: (examples_data[0][0], examples_data[0][5]), None, [active_image, active_image_path_state]
     )
+    def handle_upload(uploaded_image_np):
+        return uploaded_image_np, None, gr.update(visible=False)
+    upload_button.upload(handle_upload, upload_button, [active_image, active_image_path_state, download_button])
     generate_btn.click(
+        process_and_generate,
+        [active_image, active_image_path_state, sidebar_prompt, header_prompt, navigation_prompt, main_content_prompt],
+        [html_preview, html_code_output, download_button],
         show_progress="full"
     )
+    preview_controls = [scale_slider, width_slider, height_slider]
+    for control in preview_controls:
+        control.change(
+            render_preview,
+            [html_code_output, width_slider, height_slider, scale_slider],
+            html_preview,
+            show_progress=True
+        )
+    download_button.click(None, download_button, None, js= \
+        "(url) => { const link = document.createElement('a'); link.href = url; link.download = ''; document.body.appendChild(link); link.click(); document.body.removeChild(link); }")
+base_dir = os.path.dirname(os.path.abspath(__file__))
+allowed_paths = [
+    os.path.join(base_dir),
+    os.path.join(base_dir, 'screencoder', 'data'),
+    os.path.join(base_dir, 'screencoder', 'data', 'output'),
+    os.path.join(base_dir, 'screencoder', 'data', 'packages')
+]
+print(f"Allowed paths for file serving:")
+for path in allowed_paths:
+    print(f"  - {path}")
 if __name__ == "__main__":
+    demo.launch(allowed_paths=allowed_paths)

screencoder/block_parsor.py CHANGED Viewed

@@ -16,11 +16,6 @@ PROMPT_MERGE = "Return the bounding boxes of the sidebar, main content, header,
 BBOX_TAG_START = "<bbox>"
 BBOX_TAG_END = "</bbox>"
-# PROMPT_sidebar = "框出网页中的sidebar的位置，请你只返回对应的bounding box。"
-# PROMPT_header = "框出网页中的header的位置，请你只返回对应的bounding box。"
-# PROMPT_navigation = "框出网页中的navigation的位置，请你只返回对应的bounding box。"
-# PROMPT_main_content = "框出网页中的main content的位置，请你只返回对应的bounding box。"
 def get_args():
     parser = argparse.ArgumentParser(description="Parses bounding boxes from an image using a vision model.")
     parser.add_argument('--run_id', type=str, required=True, help='A unique identifier for the processing run.')

 BBOX_TAG_START = "<bbox>"
 BBOX_TAG_END = "</bbox>"
 def get_args():
     parser = argparse.ArgumentParser(description="Parses bounding boxes from an image using a vision model.")
     parser.add_argument('--run_id', type=str, required=True, help='A unique identifier for the processing run.')

screencoder/html_generator.py CHANGED Viewed

@@ -18,27 +18,60 @@ def get_args():
 def get_prompt_dict(instructions):
     """Dynamically creates the prompt dictionary with the user's instructions."""
     return {
-        "sidebar": f"""This is a screenshot of a container. Please fill in a complete HTML and tail-wind CSS code to accurately reproduce the given container. Please note that the layout, icon style, size, and text information of all blocks need to be basically consistent with the original screenshot based on the user's additional conditions. User instruction: {instructions["sidebar"]}. The following is the code for filling in:
-    <div>
-    your code here
-    </div>,
-    only return the code within the <div> and </div> tags""",
-        "header": f"""This is a screenshot of a container. Please fill in a complete HTML and tail-wind CSS code to accurately reproduce the given container. Please note that the relative position, layout, text information, and color of all blocks in the boundary box need to be basically consistent with the original screenshot based on the user's additional conditions. User instruction: {instructions["header"]}. The following is the code for filling in:
-    <div>
-    your code here
-    </div>,
-    only return the code within the <div> and </div> tags""",
-        "navigation": f"""This is a screenshot of a container. Please fill in a complete HTML and tail-wind CSS code to accurately reproduce the given container. Please note that the relative position, layout, text information, and color of all blocks in the boundary box need to be basically consistent with the original screenshot based on the user's additional conditions. Please use the same icons as in the original screenshot. User instruction: {instructions["navigation"]}. The following is the code for filling in:
-    <div>
-    your code here
-    </div>,
-    only return the code within the <div> and </div> tags""",
-        "main content": f"""This is a screenshot of a container. Please fill in a complete HTML and tail-wind CSS code to accurately reproduce the given container. Please note that all images displayed in the screenshot must be replaced with pure gray-400 image blocks of the same size as the corresponding images in the original screenshot, and the text information in the images does not need to be recognized. The relative position, layout, text information, and color of all blocks in the boundary box need to be basically consistent with the original screenshot based on the user's additional conditions. User instruction: {instructions["main content"]}. The following is the code for filling in:
-    <div>
-    your code here
-    </div>,
-    only return the code within the <div> and </div> tags""",
     }
 def generate_code(bbox_tree, img_path, bot, instructions):
@@ -146,6 +179,66 @@ def generate_html(bbox_tree, output_file):
     with open(output_file, 'w') as f:
         f.write(bs4.BeautifulSoup(html_content, 'html.parser').prettify())
 def code_substitution(html_file, code_dict):
     """Substitutes the generated code into the HTML file."""
@@ -212,7 +305,7 @@ def main():
         exit(1)
     bot = Doubao(api_path, model="doubao-1.5-thinking-vision-pro-250428")
-    code_dict = generate_code(root, img_path, bot, user_instruction)
     code_substitution(output_html_path, code_dict)
     print(f"HTML layout with generated content saved to {os.path.basename(output_html_path)}")

 def get_prompt_dict(instructions):
     """Dynamically creates the prompt dictionary with the user's instructions."""
+    # return {
+    #     "sidebar": f"""This is a screenshot of a container. Please fill in a complete HTML and tail-wind CSS code to accurately reproduce the given container. Please note that the layout, icon style, size, and text information of all blocks need to be basically consistent with the original screenshot based on the user's additional conditions. User instruction: {instructions["sidebar"]}. The following is the code for filling in:
+    # <div>
+    # your code here
+    # </div>,
+    # only return the code within the <div> and </div> tags""",
+    #     "header": f"""This is a screenshot of a container. Please fill in a complete HTML and tail-wind CSS code to accurately reproduce the given container. Please note that the relative position, layout, text information, and color of all blocks in the boundary box need to be basically consistent with the original screenshot based on the user's additional conditions. User instruction: {instructions["header"]}. The following is the code for filling in:
+    # <div>
+    # your code here
+    # </div>,
+    # only return the code within the <div> and </div> tags""",
+    #     "navigation": f"""This is a screenshot of a container. Please fill in a complete HTML and tail-wind CSS code to accurately reproduce the given container. Please note that the relative position, layout, text information, and color of all blocks in the boundary box need to be basically consistent with the original screenshot based on the user's additional conditions. Please use the same icons as in the original screenshot. User instruction: {instructions["navigation"]}. The following is the code for filling in:
+    # <div>
+    # your code here
+    # </div>,
+    # only return the code within the <div> and </div> tags""",
+    #     "main content": f"""This is a screenshot of a container. Please fill in a complete HTML and tail-wind CSS code to accurately reproduce the given container. Please note that all images displayed in the screenshot must be replaced with pure gray-400 image blocks of the same size as the corresponding images in the original screenshot, and the text information in the images does not need to be recognized. The relative position, layout, text information, and color of all blocks in the boundary box need to be basically consistent with the original screenshot based on the user's additional conditions. User instruction: {instructions["main content"]}. The following is the code for filling in:
+    # <div>
+    # your code here
+    # </div>,
+    # only return the code within the <div> and </div> tags""",
+    # }
     return {
+        "sidebar": f"""这是一个container的截图。这是用户给的额外要求：{instructions["sidebar"]}请填写一段完整的HTML和tail-wind CSS代码以准确再现给定的容器。请注意所有组块的排版、图标样式、大小、文字信息需要在用户额外条件的基础上与原始截图基本保持一致。请保持图标一致，如果是大块图像，请使用"placeholder.png"替换，注意代码中的大小和位置信息。以下是供填写的代码：
+        <div>
+        your code here
+        </div>
+        只需返回<div>和</div>标签内的代码""",
+        "header": f"""这是一个container的截图。这是用户给的额外要求：{instructions["header"]}请填写一段完整的HTML和tail-wind CSS代码以准确再现给定的容器。请注意所有组块在boundary box中的相对位置、排版、文字信息、颜色需要在用户额外条件的基础上与原始截图基本保持一致。以下是供填写的代码：
+        <div>
+        your code here
+        </div>
+        只需返回<div>和</div>标签内的代码""",
+        "navigation": f"""这是一个container的截图。这是用户给的额外要求：{instructions["navigation"]}请填写一段完整的HTML和tail-wind CSS代码以准确再现给定的容器。请注意所有组块的在boundary box中的相对位置、文字排版、颜色需要在用户额外条件的基础上与原始截图基本保持一致。图像请你直接使用原始截图中一致的图标。以下是供填写的代码：
+        <div>
+        your code here
+        </div>
+        只需返回<div>和</div>标签内的代码""",
+        "main content": f"""这是一个container的截图。这是用户给的额外要求：{instructions["main content"]}请填写一段完整的HTML和tail-wind CSS代码以准确再现给定的容器。截图中显示的图像务必使用"placeholder.png"替换，注意代码中其大小和位置信息，不需要识别其中的文字信息。请注意所有组块在boundary box中的相对位置、排版、文字信息、颜色需要在用户额外条件的基础上与原始截图基本保持一致。以下是供填写的代码：
+        <div>
+        your code here
+        </div>
+        只需返回<div>和</div>标签内的代码"""
     }
 def generate_code(bbox_tree, img_path, bot, instructions):
     with open(output_file, 'w') as f:
         f.write(bs4.BeautifulSoup(html_content, 'html.parser').prettify())
+def generate_code_parallel(bbox_tree, img_path, bot, instructions):
+    """generate code for all the leaf nodes in the bounding box tree, return a dictionary: {'id': 'code'}"""
+    code_dict = {}
+    t_list = []
+    prompt_dict = get_prompt_dict(instructions)
+    def _generate_code_with_retry(node, max_retries=3, retry_delay=2):
+        """Generate code with retry mechanism for rate limit errors"""
+        try:
+            # Create a new image instance for each thread
+            with Image.open(img_path) as img:
+                bbox = node["bbox"]
+                cropped_img = img.crop(bbox)
+                # Select prompt based on node type
+                if "type" in node:
+                    if node["type"] in prompt_dict:
+                        prompt = prompt_dict[node["type"]]
+                    else:
+                        print(f"Unknown component type: {node['type']}")
+                        code_dict[node["id"]] = f"<!-- Unknown component type: {node['type']} -->"
+                        return
+                else:
+                    print("Node type not found")
+                    code_dict[node["id"]] = f"<!-- Node type not found -->"
+                    return
+                for attempt in range(max_retries):
+                    try:
+                        code = bot.ask(prompt, encode_image(cropped_img))
+                        code_dict[node["id"]] = code
+                        return
+                    except Exception as e:
+                        if "rate_limit" in str(e).lower() and attempt < max_retries - 1:
+                            print(f"Rate limit hit, retrying in {retry_delay} seconds... (Attempt {attempt + 1}/{max_retries})")
+                            time.sleep(retry_delay)
+                            retry_delay *= 2  # Exponential backoff
+                        else:
+                            print(f"Error generating code for node {node['id']}: {str(e)}")
+                            code_dict[node["id"]] = f"<!-- Error: {str(e)} -->"
+                            return
+        except Exception as e:
+            print(f"Error processing image for node {node['id']}: {str(e)}")
+            code_dict[node["id"]] = f"<!-- Error: {str(e)} -->"
+    def _generate_code(node):
+        if not node.get("children"):
+            t = Thread(target=_generate_code_with_retry, args=(node,))
+            t.start()
+            t_list.append(t)
+        else:
+            for child in node["children"]:
+                _generate_code(child)
+    _generate_code(bbox_tree)
+    # Wait for all threads to complete
+    for t in t_list:
+        t.join()
+    return code_dict
 def code_substitution(html_file, code_dict):
     """Substitutes the generated code into the HTML file."""
         exit(1)
     bot = Doubao(api_path, model="doubao-1.5-thinking-vision-pro-250428")
+    code_dict = generate_code_parallel(root, img_path, bot, user_instruction)
     code_substitution(output_html_path, code_dict)
     print(f"HTML layout with generated content saved to {os.path.basename(output_html_path)}")

screencoder/image_box_detection.py CHANGED Viewed

@@ -15,18 +15,15 @@ async def extract_bboxes_from_html(html_path: Path):
         metrics = await page.evaluate("""
             () => {
-                // 1. Find and store region containers and their bboxes
                 const region_containers = Array.from(document.querySelectorAll('.box[id]'));
                 const region_bboxes = region_containers.map(el => {
                     const rect = el.getBoundingClientRect();
                     return { id: el.id, x: rect.x, y: rect.y, w: rect.width, h: rect.height };
                 });
-                // 2. Find all potential placeholders on the page
                 const placeholder_bboxes = [];
                 let ph_id_counter = 0;
-                //精准检测
-                const all_potential_placeholders = document.querySelectorAll('.bg-gray-400');
                 for (const el of all_potential_placeholders) {
                     // Apply the same filters as before
@@ -47,7 +44,6 @@ async def extract_bboxes_from_html(html_path: Path):
                         }
                     }
-                    // Only include placeholders that are inside a detected region
                     if (containing_region_id) {
                         placeholder_bboxes.push({
                             id: 'ph' + ph_id_counter++,

         metrics = await page.evaluate("""
             () => {
                 const region_containers = Array.from(document.querySelectorAll('.box[id]'));
                 const region_bboxes = region_containers.map(el => {
                     const rect = el.getBoundingClientRect();
                     return { id: el.id, x: rect.x, y: rect.y, w: rect.width, h: rect.height };
                 });
                 const placeholder_bboxes = [];
                 let ph_id_counter = 0;
+                const all_potential_placeholders = document.querySelectorAll('img[src="placeholder.png"]');
                 for (const el of all_potential_placeholders) {
                     // Apply the same filters as before
                         }
                     }
                     if (containing_region_id) {
                         placeholder_bboxes.push({
                             id: 'ph' + ph_id_counter++,

screencoder/image_replacer.py CHANGED Viewed

@@ -99,7 +99,7 @@ def main():
     soup = BeautifulSoup(html_content, 'html.parser')
     # 1. Find all placeholder elements by their class, in document order.
-    placeholder_elements = soup.find_all(class_="bg-gray-400")
     # 2. Get the placeholder IDs from the mapping file in the correct, sorted order.
     def natural_sort_key(s):
@@ -115,32 +115,27 @@ def main():
     # 3. Check for count mismatches
     if len(placeholder_elements) != len(ordered_placeholder_ids):
-        print(f"Warning: Mismatch in counts! Found {len(placeholder_elements)} gray boxes in HTML, but {len(ordered_placeholder_ids)} mappings.")
     else:
-        print(f"Found {len(placeholder_elements)} gray boxes to replace.")
     # 4. Iterate through both lists, create a proper <img> tag, and replace the placeholder.
     for i, ph_element in enumerate(placeholder_elements):
         if i >= len(ordered_placeholder_ids):
-            print(f"Warning: More gray boxes in HTML than mappings. Stopping at box {i+1}.")
             break
         ph_id = ordered_placeholder_ids[i]
         relative_img_path = f"{crop_dir.name}/{ph_id}.png"
-        # --- Create a new <img> tag and replace the placeholder ---
-        # a. Get all classes from the original placeholder to preserve styling.
-        original_classes = ph_element.get('class', [])
-        if 'bg-gray-400' in original_classes:
-            original_classes.remove('bg-gray-400') # Remove the placeholder background
-        # b. Create the new <img> tag
-        img_tag = soup.new_tag("img", src=relative_img_path)
-        img_tag['class'] = original_classes
-        # c. Replace the placeholder with the new image tag.
-        ph_element.replace_with(img_tag)
     # Save the modified HTML
     final_html_path.write_text(str(soup))

     soup = BeautifulSoup(html_content, 'html.parser')
     # 1. Find all placeholder elements by their class, in document order.
+    placeholder_elements = soup.find_all('img', src="placeholder.png")
     # 2. Get the placeholder IDs from the mapping file in the correct, sorted order.
     def natural_sort_key(s):
     # 3. Check for count mismatches
     if len(placeholder_elements) != len(ordered_placeholder_ids):
+        print(f"Warning: Mismatch in counts! Found {len(placeholder_elements)} placeholder images in HTML, but {len(ordered_placeholder_ids)} mappings.")
     else:
+        print(f"Found {len(placeholder_elements)} placeholder images to replace.")
     # 4. Iterate through both lists, create a proper <img> tag, and replace the placeholder.
     for i, ph_element in enumerate(placeholder_elements):
         if i >= len(ordered_placeholder_ids):
+            print(f"Warning: More placeholder images in HTML than mappings. Stopping at image {i+1}.")
             break
         ph_id = ordered_placeholder_ids[i]
+        # Fix: Use the correct relative path from HTML file to image directory
         relative_img_path = f"{crop_dir.name}/{ph_id}.png"
+        # Debug: Print the path being used
+        print(f"Setting image path for {ph_id}: {relative_img_path}")
+        # --- Update the img tag's src attribute ---
+        # Since we're now working with img tags instead of div tags,
+        # we just need to update the src attribute
+        ph_element['src'] = relative_img_path
     # Save the modified HTML
     final_html_path.write_text(str(soup))

screencoder/main.py CHANGED Viewed

@@ -99,13 +99,15 @@ def generate_html_for_demo(image_path, instructions):
             with open(final_html_path, 'r', encoding='utf-8') as f:
                 html_content = f.read()
             print(f"Successfully generated HTML for run_id: {run_id}")
-            return html_content
         else:
-            return f"Error: Final HTML file not found for run_id: {run_id}"
     except Exception as e:
         print(f"An error occurred during the workflow for run_id {run_id}: {e}")
-        return f"An error occurred: {e}"
     finally:
         # 4. Cleanup: Remove temporary directories
         try:

             with open(final_html_path, 'r', encoding='utf-8') as f:
                 html_content = f.read()
             print(f"Successfully generated HTML for run_id: {run_id}")
+            return html_content, run_id
         else:
+            error_msg = f"Error: Final HTML file not found for run_id: {run_id}"
+            return error_msg, run_id
     except Exception as e:
+        error_msg = f"An error occurred: {e}"
         print(f"An error occurred during the workflow for run_id {run_id}: {e}")
+        return error_msg, run_id
     finally:
         # 4. Cleanup: Remove temporary directories
         try:

screencoder/mapping.py CHANGED Viewed

@@ -319,8 +319,6 @@ def main():
     # Always generate the debug image if the source exists
     generate_debug_overlay(debug_src_path, all_uied_boxes, final_results, uied_shape, debug_overlay_path)
-    print(f"Debug image written to {debug_overlay_path}")
     print(f"--- Mapping Complete for run_id: {run_id} ---")
 def get_args():

     # Always generate the debug image if the source exists
     generate_debug_overlay(debug_src_path, all_uied_boxes, final_results, uied_shape, debug_overlay_path)
     print(f"--- Mapping Complete for run_id: {run_id} ---")
 def get_args():