screenshot2html

Runtime error

File size: 7,109 Bytes

import os
import subprocess

from playwright.sync_api import sync_playwright
from typing import List
from PIL import Image

import gradio as gr
from gradio_client.client import DEFAULT_TEMP_DIR
from transformers import AutoProcessor, AutoModelForCausalLM


API_TOKEN = os.getenv("HF_AUTH_TOKEN")
# PROCESSOR = AutoProcessor.from_pretrained(
#     "HuggingFaceM4/img2html",
#     token=API_TOKEN,
# )
# MODEL = AutoModelForCausalLM.from_pretrained(
#     "HuggingFaceM4/img2html",
#     token=API_TOKEN,
# )


IMAGE_GALLERY_PATHS = [
    f"example_images/{ex_image}"
    for ex_image in os.listdir(f"example_images")
]


def install_playwright():
    try:
        subprocess.run(["playwright", "install"], check=True)
        print("Playwright installation successful.")
    except subprocess.CalledProcessError as e:
        print(f"Error during Playwright installation: {e}")

install_playwright()

def add_file_gallery(
    selected_state: gr.SelectData,
    gallery_list: List[str]
):
    return f"example_images/{gallery_list.root[selected_state.index].image.orig_name}"

def render_webpage(
    html_css_code,
):
    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        context = browser.new_context(
            user_agent=(
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/107.0.0.0"
                " Safari/537.36"
            )
        )
        page = context.new_page()
        page.set_content(html_css_code)
        page.wait_for_load_state("networkidle")
        output_path_screenshot = f"{DEFAULT_TEMP_DIR}/{hash(html_css_code)}.png"
        _ = page.screenshot(path=output_path_screenshot, full_page=True)

        context.close()
        browser.close()

    return Image.open(output_path_screenshot)


def model_inference(
    image,
):
    CAR_COMPNAY = """<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>XYZ Car Company</title>
    <style>
        body {
            font-family: 'Arial', sans-serif;
            margin: 0;
            padding: 0;
            background-color: #f4f4f4;
        }

        header {
            background-color: #333;
            color: #fff;
            padding: 1em;
            text-align: center;
        }

        nav {
            background-color: #555;
            color: #fff;
            padding: 0.5em;
            text-align: center;
        }

        nav a {
            color: #fff;
            text-decoration: none;
            padding: 0.5em 1em;
            margin: 0 1em;
        }

        section {
            padding: 2em;
        }

        h2 {
            color: #333;
        }

        .car-container {
            display: flex;
            flex-wrap: wrap;
            justify-content: space-around;
        }

        .car-card {
            width: 300px;
            margin: 1em;
            border: 1px solid #ddd;
            border-radius: 5px;
            overflow: hidden;
            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
        }

        .car-image {
            width: 100%;
            height: 150px;
            object-fit: cover;
        }

        .car-details {
            padding: 1em;
        }

        footer {
            background-color: #333;
            color: #fff;
            text-align: center;
            padding: 1em;
            position: fixed;
            bottom: 0;
            width: 100%;
        }
    </style>
</head>
<body>

    <header>
        <h1>XYZ Car Company</h1>
    </header>

    <nav>
        <a href="#">Home</a>
        <a href="#">Models</a>
        <a href="#">About Us</a>
        <a href="#">Contact</a>
    </nav>

    <section>
        <h2>Our Cars</h2>
        <div class="car-container">
            <div class="car-card">
                <img src="car1.jpg" alt="Car 1" class="car-image">
                <div class="car-details">
                    <h3>Model A</h3>
                    <p>Description of Model A.</p>
                </div>
            </div>

            <div class="car-card">
                <img src="car2.jpg" alt="Car 2" class="car-image">
                <div class="car-details">
                    <h3>Model B</h3>
                    <p>Description of Model B.</p>
                </div>
            </div>

            <!-- Add more car cards as needed -->
        </div>
    </section>

    <footer>
        &copy; 2024 XYZ Car Company. All rights reserved.
    </footer>

</body>
</html>"""
    rendered_page = render_webpage(CAR_COMPNAY)
    return CAR_COMPNAY, rendered_page


generated_html = gr.Code(
    label="Extracted HTML",
    elem_id="generated_html",
)
rendered_html = gr.Image(
    label="Rendered HTML"
)
# rendered_html = gr.HTML(
#     label="Rendered HTML"
# )


css = """
.gradio-container{max-width: 1000px!important}
h1{display: flex;align-items: center;justify-content: center;gap: .25em}
*{transition: width 0.5s ease, flex-grow 0.5s ease}
"""


with gr.Blocks(title="Img2html", theme=gr.themes.Base(), css=css) as demo:
    with gr.Row(equal_height=True):
        with gr.Column(scale=4, min_width=250) as upload_area:
            imagebox = gr.Image(
                type="filepath",
                label="Screenshot to extract",
                visible=True,
                sources=["upload", "clipboard"],
            )
            with gr.Group():
                with gr.Row():
                    submit_btn = gr.Button(
                        value="▶️ Submit", visible=True, min_width=120
                    )
                    clear_btn = gr.ClearButton(
                        [imagebox, generated_html, rendered_html], value="🧹 Clear", min_width=120
                    )
                    regenerate_btn = gr.Button(
                        value="🔄 Regenerate", visible=True, min_width=120
                    )
        with gr.Column(scale=4) as result_area:
            rendered_html.render()

    with gr.Row():
        generated_html.render()

    with gr.Row():
        template_gallery = gr.Gallery(
            value=IMAGE_GALLERY_PATHS,
            label="Templates Gallery",
            allow_preview=False,
            columns=4,
            elem_id="gallery",
            show_share_button=False,
            height=400,
        )

    gr.on(
        triggers=[
            imagebox.upload,
            submit_btn.click,
            template_gallery.select,
            regenerate_btn.click,
        ],
        fn=model_inference,
        inputs=[imagebox],
        outputs=[generated_html, rendered_html],
        queue=False,
    )
    regenerate_btn.click(
        fn=model_inference,
        inputs=[
            imagebox,
        ],
        outputs=[generated_html, rendered_html],
        queue=False,
    )
    template_gallery.select(
        fn=add_file_gallery,
        inputs=[template_gallery],
        outputs=[imagebox],
        queue=False,
    )
    demo.load(queue=False)

demo.queue(max_size=40, api_open=False)
demo.launch(max_threads=400)