File size: 1,935 Bytes
4da1734 | 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 | import os
import re
from io import BytesIO
from PIL import Image
from playwright.sync_api import sync_playwright
def extract_clean_html(text):
"""
清理模型输出，尽量提取完整 HTML。
"""
text = text.replace("```html", "").replace("```", "")
start_match = re.search(r"<!DOCTYPE html>", text, re.IGNORECASE)
end_match = re.search(r"</html>", text, re.IGNORECASE)
if start_match and end_match:
start_idx = start_match.start()
end_idx = end_match.end()
if end_idx > start_idx:
return text[start_idx:end_idx]
return text.strip()
def render_html_to_image(html, width=1080, height=2400):
"""
用 Playwright 把 HTML 渲染成 PIL.Image。
"""
with sync_playwright() as p:
browser = p.chromium.launch(headless=True)
page = browser.new_page(viewport={"width": width, "height": height})
page.set_content(html, wait_until="domcontentloaded")
screenshot_bytes = page.screenshot(full_page=False)
browser.close()
return Image.open(BytesIO(screenshot_bytes)).convert("RGB")
def save_demo_outputs(output_dir, hinted_image, html, rendered_image=None):
"""
保存 demo 输出：
- input_with_hint.png
- predicted_next_ui.html
- predicted_next_ui.png
"""
os.makedirs(output_dir, exist_ok=True)
hinted_path = os.path.join(output_dir, "input_with_hint.png")
html_path = os.path.join(output_dir, "predicted_next_ui.html")
rendered_path = os.path.join(output_dir, "predicted_next_ui.png")
hinted_image.save(hinted_path)
with open(html_path, "w", encoding="utf-8") as f:
f.write(html)
if rendered_image is not None:
rendered_image.save(rendered_path)
print(f"Saved hinted image to: {hinted_path}")
print(f"Saved HTML to: {html_path}")
if rendered_image is not None:
print(f"Saved rendered image to: {rendered_path}") |