app.py

import gradio as gr
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from PIL import Image
from io import BytesIO
import tempfile
import time
import os

def get_full_dimensions(driver):
    """
    body.scrollWidth / scrollHeight と
    documentElement.scrollWidth / scrollHeight を比較して
    ページ全体の幅・高さを返すユーティリティ関数
    """
    body_w = driver.execute_script("return document.body.scrollWidth")
    body_h = driver.execute_script("return document.body.scrollHeight")
    html_w = driver.execute_script("return document.documentElement.scrollWidth")
    html_h = driver.execute_script("return document.documentElement.scrollHeight")

    full_width = max(body_w, html_w)
    full_height = max(body_h, html_h)
    return full_width, full_height

def render_fullpage_screenshot(html_code):
    # HTMLコードを一時ファイルに保存
    tmp_file = tempfile.NamedTemporaryFile(suffix=".html", delete=False)
    tmp_path = tmp_file.name
    tmp_file.write(html_code.encode('utf-8'))
    tmp_file.close()

    # ヘッドレスChrome起動オプション
    options = Options()
    options.add_argument("--headless")
    options.add_argument("--no-sandbox")
    options.add_argument("--disable-dev-shm-usage")
    options.add_argument("--force-device-scale-factor=1")

    try:
        driver = webdriver.Chrome(options=options)

        # 1) とりあえず適当なウィンドウサイズでページを読み込む
        driver.set_window_size(1200, 800)
        driver.get("file://" + tmp_path)

        # ページロード完了を待つ
        WebDriverWait(driver, 10).until(
            EC.presence_of_element_located((By.TAG_NAME, "body"))
        )
        time.sleep(1)  # 追加待機

        # 2) スクロールバーを非表示にする（必要に応じてコメントアウト）
        driver.execute_script("""
            document.documentElement.style.overflow = 'hidden';
            document.body.style.overflow = 'hidden';
        """)

        # 3) ページ全体のサイズを取得してウィンドウサイズを拡張
        full_width, full_height = get_full_dimensions(driver)
        driver.set_window_size(full_width, full_height)
        time.sleep(2)  # レイアウト変化待ち

        # 4) もう一度ページ全体のサイズを取得し、もし変化していれば再調整
        new_full_width, new_full_height = get_full_dimensions(driver)
        if (new_full_width != full_width) or (new_full_height != full_height):
            driver.set_window_size(new_full_width, new_full_height)
            time.sleep(2)  # もう一度待機

        # 5) ページ全体をスクリーンショット
        png = driver.get_screenshot_as_png()

    except Exception as e:
        # エラー時は1x1の黒画像を返す
        return Image.new('RGB', (1, 1), color=(0, 0, 0))

    finally:
        driver.quit()
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

    return Image.open(BytesIO(png))

# Gradioインターフェース
iface = gr.Interface(
    fn=render_fullpage_screenshot,
    inputs=gr.Textbox(lines=15, label="HTMLコード入力"),
    outputs=gr.Image(type="pil", label="ページ全体のスクリーンショット"),
    title="Full Page Screenshot App",
    description="HTMLをヘッドレスブラウザでレンダリングし、ページ全体を1枚の画像として取得します。"
)

if __name__ == "__main__":
    iface.launch()