screenshot2html / app.py
VictorSanh's picture
first commit
1896bc8
raw
history blame
7.23 kB
from typing import List, Optional, Tuple
from PIL import Image
from playwright.sync_api import sync_playwright
import os
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,
# )
IMAGE_GALLERY_PATHS = [
f"example_images/{ex_image}"
for ex_image in os.listdir(f"example_images")
]
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}",
# "",
# )
return f"example_images/{gallery_list.root[selected_state.index].image.orig_name}"
def expand_layout():
return gr.Column(scale=2), gr.Textbox()
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
# textbox = gr.Textbox(
# placeholder="Upload an image and ask the AI to create a meme!",
# show_label=False,
# value="Write a meme about this image.",
# visible=True,
# container=False,
# label="Text input",
# scale=8,
# max_lines=5,
# )
generated_html = gr.Textbox(
label="IDEFICS Generated HTML",
elem_id="generated_html",
)
rendered_html = gr.Image(
)
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):
# scale=2 when expanded
with gr.Column(scale=4, min_width=250) as upload_area:
imagebox = gr.Image(
type="filepath", label="Image to HTML", height=272, 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=5) as result_area:
rendered_html.render()
with gr.Row():
generated_html.render()
with gr.Row(equal_height=True):
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(
# fn=choose_gallery,
# inputs=[gallery_type_choice],
# outputs=[template_gallery],
queue=False,
)
demo.queue(max_size=40, api_open=False)
demo.launch(max_threads=400)