Update app.py

app.py

@@ -1,3 +1,5 @@
+# --- START OF FILE app.py ---
+
 import gradio as gr
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
@@ -9,87 +11,197 @@ from io import BytesIO
 import tempfile
 import time
 import os
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+import logging
+
+# ロギング設定 (デバッグに役立つ)
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-def render_fullpage_screenshot(html_code, extension_percentage):
+# --- コアロジック: スクリーンショット生成関数 ---
+# (元のコードから変更なし、エラーハンドリングを少し明確化)
+def render_fullpage_screenshot(html_code: str, extension_percentage: float) -> Image.Image:
     """
     html_code: HTMLのソースコード文字列
     extension_percentage: 上下に追加する余裕の% (例: 4なら合計4%の余裕)
+    戻り値: PIL Imageオブジェクト。エラー時は1x1の黒画像。
     """
-    # 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")
+    tmp_file = None
+    driver = None
+    tmp_path = ""
 
     try:
+        # 1) HTMLコードを一時ファイルに保存
+        # delete=False と NamedTemporaryFile の close/remove の組み合わせはLinux/Windowsで挙動が安定しないことがあるため、
+        # with構文と手動削除を確実に組み合わせる
+        with tempfile.NamedTemporaryFile(suffix=".html", delete=False, mode='w', encoding='utf-8') as tmp_file:
+            tmp_path = tmp_file.name
+            tmp_file.write(html_code)
+        logger.info(f"HTML saved to temporary file: {tmp_path}")
+
+        # 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") # DPIスケーリング無効化
+        options.add_argument("--window-size=1200,800") # 初期ウィンドウサイズ
+        # デバッグ用にログレベルを上げる (オプション)
+        # options.add_argument("--enable-logging")
+        # options.add_argument("--v=1")
+
+        logger.info("Initializing WebDriver...")
+        # サービスオブジェクトを作成してログ出力を制御 (オプション)
+        # service = webdriver.chrome.service.Service(log_output=subprocess.STDOUT)
+        # driver = webdriver.Chrome(options=options, service=service)
         driver = webdriver.Chrome(options=options)
+        logger.info("WebDriver initialized.")
 
-        # 3) 初期ウィンドウサイズでページを読み込む
-        driver.set_window_size(1200, 800)
-        driver.get("file://" + tmp_path)
+        # 3) 初期ウィンドウサイズでページを読み込む (オプションで設定済みなので不要かも)
+        # driver.set_window_size(1200, 800)
+        file_url = "file://" + tmp_path
+        logger.info(f"Navigating to {file_url}")
+        driver.get(file_url)
 
-        # 4) ページのロードを待機
-        WebDriverWait(driver, 10).until(
+        # 4) ページのロードを待機 (body要素の存在確認)
+        logger.info("Waiting for body element...")
+        WebDriverWait(driver, 20).until( # タイムアウトを少し延長
             EC.presence_of_element_located((By.TAG_NAME, "body"))
         )
-        time.sleep(2)  # 外部リソースの読み込み待機
+        logger.info("Body element found. Waiting for potential dynamic content loading...")
+        time.sleep(3) # 外部リソースやJSによる動的変更の待機時間を少し増やす
 
         # 5) スクロールバーがキャプチャに写らないようCSSで非表示
+        logger.info("Hiding scrollbars via JS...")
         driver.execute_script(
             "document.documentElement.style.overflow = 'hidden';"
             "document.body.style.overflow = 'hidden';"
         )
 
-        # 6) ページ全体の幅・高さを正確に取得 (bodyとdocumentElementの両方から)
+        # 6) ページ全体の幅・高さを正確に取得
+        logger.info("Calculating page dimensions...")
         scroll_width = driver.execute_script(
-            "return Math.max(document.body.scrollWidth, document.documentElement.scrollWidth)"
+            "return Math.max(document.body.scrollWidth, document.documentElement.scrollWidth, document.body.offsetWidth, document.documentElement.offsetWidth)"
         )
         scroll_height = driver.execute_script(
-            "return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight)"
+            "return Math.max(document.body.scrollHeight, document.documentElement.scrollHeight, document.body.offsetHeight, document.documentElement.offsetHeight)"
         )
+        logger.info(f"Calculated dimensions: width={scroll_width}, height={scroll_height}")
+
+        # 幅が0や極端に小さい場合はデフォルト値を使うなど対策が必要になる可能性あり
+        if scroll_width < 100:
+           logger.warning(f"Detected very small scroll_width ({scroll_width}), defaulting to 1200")
+           scroll_width = 1200
 
         # 7) 上下方向にだけユーザー指定の余裕(%)を加えた高さを計算
-        adjusted_height = int(scroll_height * (1 + extension_percentage/100))
-        
+        adjusted_height = int(scroll_height * (1 + extension_percentage / 100))
+        logger.info(f"Adjusted height: {adjusted_height} (extension: {extension_percentage}%)")
+
         # 8) ウィンドウサイズを、幅はそのまま、縦はadjusted_heightに変更
+        logger.info(f"Resizing window to {scroll_width} x {adjusted_height}")
+        # set_window_sizeが期待通りに動作しない場合があるため、最大化してからサイズ変更を試すことも検討
+        # driver.maximize_window()
         driver.set_window_size(scroll_width, adjusted_height)
-        time.sleep(2)  # レイアウトの安定化待機
+        logger.info("Waiting for layout stabilization after resize...")
+        time.sleep(3) # レイアウトの安定化待機時間を少し増やす
 
         # 念のため最上部にスクロール
+        logger.info("Scrolling to top...")
         driver.execute_script("window.scrollTo(0, 0)")
         time.sleep(1)
 
         # 9) スクリーンショット取得
+        logger.info("Taking screenshot...")
         png = driver.get_screenshot_as_png()
+        logger.info("Screenshot taken successfully.")
+
+        return Image.open(BytesIO(png))
 
     except Exception as e:
+        logger.error(f"Error during screenshot generation: {e}", exc_info=True)
         # エラー時は1x1の黒画像を返す
         return Image.new('RGB', (1, 1), color=(0, 0, 0))
     finally:
-        driver.quit()
-        if os.path.exists(tmp_path):
-            os.remove(tmp_path)
+        # 必ず WebDriver を終了する
+        if driver:
+            logger.info("Quitting WebDriver...")
+            driver.quit()
+            logger.info("WebDriver quit.")
+        # 一時ファイルを確実に削除する
+        if tmp_path and os.path.exists(tmp_path):
+            try:
+                os.remove(tmp_path)
+                logger.info(f"Temporary file removed: {tmp_path}")
+            except OSError as e:
+                logger.error(f"Error removing temporary file {tmp_path}: {e}")
+
+# --- FastAPI アプリケーション設定 ---
+app = FastAPI(
+    title="Full Page Screenshot API",
+    description="Renders HTML using a headless browser and returns a full-page screenshot.",
+    version="1.0.0"
+)
+
+# --- APIリクエストモデル ---
+class ScreenshotRequest(BaseModel):
+    html_code: str = Field(..., description="HTML source code to render.")
+    extension_percentage: float = Field(8.0, ge=0, le=100, description="Percentage of extra height to add (top and bottom combined). Default: 8%")
 
-    return Image.open(BytesIO(png))
+# --- APIエンドポイント ---
+@app.post("/screenshot",
+          response_class=StreamingResponse,
+          tags=["Screenshot"],
+          summary="Generate Full Page Screenshot",
+          description="Takes HTML code and returns a PNG image of the rendered full page.")
+async def create_screenshot(request: ScreenshotRequest):
+    """
+    Generates a full-page screenshot from the provided HTML code.
+    """
+    logger.info(f"Received API request. HTML length: {len(request.html_code)}, Extension: {request.extension_percentage}%")
+    try:
+        pil_image = render_fullpage_screenshot(request.html_code, request.extension_percentage)
+
+        # エラー画像 (1x1) が返ってきたかチェック
+        if pil_image.size == (1, 1):
+             logger.error("Screenshot generation failed, returning 500 error.")
+             # エラーを示すために 500 Internal Server Error を返すことも検討できる
+             # raise HTTPException(status_code=500, detail="Failed to generate screenshot. Check HTML or server logs.")
+             # ここでは仕様通り 1x1 画像を返すことにするが、APIとしてはエラーの方が親切かもしれない
+
+        # PIL画像をPNGバイトデータに変換
+        img_byte_arr = BytesIO()
+        pil_image.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0) # ポインタを先頭に戻す
 
-# Gradioインターフェースの定義
-iface = gr.Interface(
-    fn=render_fullpage_screenshot,
+        logger.info("Returning screenshot as PNG stream.")
+        return StreamingResponse(img_byte_arr, media_type="image/png")
+
+    except Exception as e:
+        logger.error(f"Unhandled exception in API endpoint: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
+
+# --- Gradioインターフェース定義 ---
+gradio_interface = gr.Interface(
+    fn=render_fullpage_screenshot, # コアロジック関数を再利用
     inputs=[
         gr.Textbox(lines=15, label="HTMLコード入力"),
-        gr.Slider(minimum=0, maximum=20, step=1.0, value=8, label="上下高さ拡張率（%）")
+        gr.Slider(minimum=0, maximum=30, step=1.0, value=8, label="上下高さ拡張率（%）") # 最大値を少し増やす
     ],
     outputs=gr.Image(type="pil", label="ページ全体のスクリーンショット"),
-    title="Full Page Screenshot API (高さ拡張調整可能)",
-    description="HTMLをヘッドレスブラウザでレンダリングし、ページ全体を1枚の画像として取得します。上下のみユーザー指定の余裕(%)を追加します。"
+    title="Full Page Screenshot (高さ拡張調整可能)",
+    description="HTMLをヘッドレスブラウザでレンダリングし、ページ全体を1枚の画像として取得します。上下のみユーザー指定の余裕(%)を追加します。APIは /screenshot (POST) で利用可能です。",
+    allow_flagging="never" # Hugging Face Spacesでフラグ機能を無効化
 )
 
-if __name__ == "__main__":
-    iface.launch()
+# --- GradioアプリをFastAPIにマウント ---
+# ルートパス ("/") でGradioインターフェースを提供
+app = gr.mount_gradio_app(app, gradio_interface, path="/")
+
+# --- uvicorn で実行するための設定 (Hugging Face Spacesでは通常不要) ---
+# if __name__ == "__main__":
+#     import uvicorn
+#     uvicorn.run(app, host="0.0.0.0", port=7860)
+
+# --- END OF FILE app.py ---