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 render_fullpage_screenshot_with_margin(html_code):
    # 1) 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()

    # 2) ヘッドレスChrome(Chromium)起動オプションを設定
    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)

        # 3) 初期ウィンドウサイズを設定してページを読み込む
        driver.set_window_size(1200, 800)
        driver.get("file://" + tmp_path)

        # 4) ページのロード完了を待機（bodyタグの存在）
        WebDriverWait(driver, 10).until(
            EC.presence_of_element_located((By.TAG_NAME, "body"))
        )
        time.sleep(2)  # 外部リソースの読み込みを待つ

        # 5) スクロールバー非表示と背景色の固定（マルチライン文字列で記述）
        driver.execute_script(
            """
            document.documentElement.style.overflow = 'hidden';
            document.body.style.overflow = 'hidden';
            document.documentElement.style.backgroundColor = 'white';
            document.body.style.backgroundColor = 'white';
            """
        )

        # 6) ページ全体の幅と高さを、bodyとdocumentElementから取得（大きい方を使用）
        scroll_width = driver.execute_script(
            "return Math.max(document.body.scrollWidth, document.documentElement.scrollWidth);"
        )
        scroll_height = driver.execute_script(
            "return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight);"
        )

        # 7) 上下左右に5%余白を確保（合計10%拡大）
        margin_factor = 0.05  # 5%
        final_width = int(scroll_width * (1 + margin_factor * 2))
        final_height = int(scroll_height * (1 + margin_factor * 2))

        # 8) 余白込みのウィンドウサイズに再設定
        driver.set_window_size(final_width, final_height)
        time.sleep(2)  # レイアウト変化の安定待ち

        # 9) コンテンツを中央に配置するため、余白分だけスクロール
        offset_x = int(scroll_width * margin_factor)
        offset_y = int(scroll_height * margin_factor)
        driver.execute_script(f"window.scrollTo({offset_x}, {offset_y});")
        time.sleep(1)

        # 10) スクリーンショットをPNGバイナリとして取得
        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)

    # 11) PNGバイナリをPIL.Imageに変換して返す
    return Image.open(BytesIO(png))

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

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