Update app.py

app.py

@@ -1,6 +1,9 @@
 # --- START OF FILE app.py ---
 
 import gradio as gr
+from fastapi import FastAPI, HTTPException, Body
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.by import By
@@ -11,197 +14,233 @@ 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
+import logging # loggingを追加
 
-# ロギング設定 (デバッグに役立つ)
+# ロギング設定
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# --- コアロジック: スクリーンショット生成関数 ---
-# (元のコードから変更なし、エラーハンドリングを少し明確化)
+# --- Core Screenshot Logic ---
 def render_fullpage_screenshot(html_code: str, extension_percentage: float) -> Image.Image:
     """
-    html_code: HTMLのソースコード文字列
-    extension_percentage: 上下に追加する余裕の% (例: 4なら合計4%の余裕)
-    戻り値: PIL Imageオブジェクト。エラー時は1x1の黒画像。
+    Renders HTML code to a full-page screenshot using Selenium.
+
+    Args:
+        html_code: The HTML source code string.
+        extension_percentage: Percentage of extra space to add vertically (e.g., 4 means 4% total).
+
+    Returns:
+        A PIL Image object of the screenshot. Returns a 1x1 black image on error.
     """
-    tmp_file = None
-    driver = None
-    tmp_path = ""
+    tmp_path = None # 初期化
+    driver = None # 初期化
 
+    # 1) Save HTML code to a temporary file
     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}")
+    except Exception as e:
+        logger.error(f"Error writing temporary HTML file: {e}")
+        return Image.new('RGB', (1, 1), color=(0, 0, 0)) # エラー時は黒画像
+
+    # 2) Headless Chrome(Chromium) options
+    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")
+    # Increase logging verbosity for debugging if needed
+    # options.add_argument("--enable-logging")
+    # options.add_argument("--v=1")
 
-        # 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")
-
+    try:
         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)
+        # 3) Load page with initial window size
+        driver.set_window_size(1200, 800)
         file_url = "file://" + tmp_path
         logger.info(f"Navigating to {file_url}")
         driver.get(file_url)
 
-        # 4) ページのロードを待機 (body要素の存在確認)
+        # 4) Wait for page load
         logger.info("Waiting for body element...")
-        WebDriverWait(driver, 20).until( # タイムアウトを少し延長
+        WebDriverWait(driver, 15).until( # タイムアウトを少し延長
             EC.presence_of_element_located((By.TAG_NAME, "body"))
         )
-        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) ページ全体の幅・高さを正確に取得
-        logger.info("Calculating page dimensions...")
-        scroll_width = driver.execute_script(
-            "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, 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))
-        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()
+        logger.info("Body element found. Waiting for potential resource loading...")
+        time.sleep(3) # Wait a bit longer for external resources/scripts
+
+        # 5) Hide scrollbars via CSS
+        try:
+            driver.execute_script(
+                "document.documentElement.style.overflow = 'hidden';"
+                "document.body.style.overflow = 'hidden';"
+            )
+            logger.info("Scrollbars hidden via JS.")
+        except Exception as e:
+            logger.warning(f"Could not hide scrollbars via JS: {e}")
+
+
+        # 6) Get full page dimensions accurately
+        try:
+            scroll_width = driver.execute_script(
+                "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, document.body.offsetHeight, document.documentElement.offsetHeight)"
+            )
+            logger.info(f"Detected dimensions: width={scroll_width}, height={scroll_height}")
+            # Ensure minimum dimensions to avoid errors
+            scroll_width = max(scroll_width, 100) # 最小幅を設定
+            scroll_height = max(scroll_height, 100) # 最小高さを設定
+
+        except Exception as e:
+            logger.error(f"Error getting page dimensions: {e}")
+            # フォールバックとしてデフォルト値を設定
+            scroll_width = 1200
+            scroll_height = 800
+            logger.warning(f"Falling back to dimensions: width={scroll_width}, height={scroll_height}")
+
+
+        # 7) Calculate adjusted height with user-specified margin
+        adjusted_height = int(scroll_height * (1 + extension_percentage / 100.0))
+        # Ensure adjusted height is not excessively large or small
+        adjusted_height = max(adjusted_height, scroll_height, 100) # 最小高さを確保
+        logger.info(f"Adjusted height calculated: {adjusted_height} (extension: {extension_percentage}%)")
+
+
+        # 8) Set window size to full page dimensions (width) and adjusted height
+        logger.info(f"Resizing window to: width={scroll_width}, height={adjusted_height}")
         driver.set_window_size(scroll_width, adjusted_height)
         logger.info("Waiting for layout stabilization after resize...")
-        time.sleep(3) # レイアウトの安定化待機時間を少し増やす
+        time.sleep(3) # Wait longer for layout stabilization
 
-        # 念のため最上部にスクロール
-        logger.info("Scrolling to top...")
-        driver.execute_script("window.scrollTo(0, 0)")
-        time.sleep(1)
+        # Scroll to top just in case
+        try:
+            driver.execute_script("window.scrollTo(0, 0)")
+            time.sleep(1)
+            logger.info("Scrolled to top.")
+        except Exception as e:
+             logger.warning(f"Could not scroll to top: {e}")
 
-        # 9) スクリーンショット取得
+        # 9) Take screenshot
         logger.info("Taking screenshot...")
         png = driver.get_screenshot_as_png()
         logger.info("Screenshot taken successfully.")
 
-        return Image.open(BytesIO(png))
+        # Convert to PIL Image
+        img = Image.open(BytesIO(png))
+        # Crop the image back to the original scroll_height plus margin
+        # This removes extra blank space at the bottom if window resize was much larger
+        # The screenshot captures the viewport, which we set to adjusted_height
+        # We only need scroll_height + (scroll_height * extension_percentage / 200) top/bottom margin
+        # However, simple approach: use adjusted_height directly or crop if needed.
+        # Let's return the full adjusted_height capture for now.
+        # If cropping is needed:
+        # final_height = int(scroll_height * (1 + extension_percentage / 100.0))
+        # if img.height > final_height:
+        #    img = img.crop((0, 0, img.width, final_height))
+        #    logger.info(f"Cropped image to final height: {final_height}")
+
+        return img
 
     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))
+        logger.error(f"An error occurred during screenshot generation: {e}", exc_info=True)
+        # Optionally capture a screenshot even on error for debugging
+        # try:
+        #     if driver: driver.save_screenshot("error_screenshot.png")
+        # except Exception as screen_err:
+        #     logger.error(f"Could not save error screenshot: {screen_err}")
+        return Image.new('RGB', (1, 1), color=(0, 0, 0)) # Return black 1x1 image on error
     finally:
-        # 必ず WebDriver を終了する
+        logger.info("Cleaning up...")
         if driver:
-            logger.info("Quitting WebDriver...")
-            driver.quit()
-            logger.info("WebDriver quit.")
-        # 一時ファイルを確実に削除する
+            try:
+                driver.quit()
+                logger.info("WebDriver quit successfully.")
+            except Exception as e:
+                logger.error(f"Error quitting WebDriver: {e}")
         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.info(f"Temporary file {tmp_path} removed.")
+            except Exception 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"
-)
+# --- FastAPI Setup ---
+app = FastAPI()
 
-# --- APIリクエストモデル ---
+# Pydantic model for API request body validation
 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%")
+    html_code: str
+    extension_percentage: float = 8.0 # Default value same as Gradio slider
 
-# --- APIエンドポイント ---
-@app.post("/screenshot",
+# API Endpoint for screenshot generation
+@app.post("/api/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):
+          summary="Render HTML to Full Page Screenshot",
+          description="Takes HTML code and an optional vertical extension percentage, renders it using a headless browser, and returns the full-page screenshot as a PNG image.")
+async def api_render_screenshot(request: ScreenshotRequest):
     """
-    Generates a full-page screenshot from the provided HTML code.
+    API endpoint to render HTML and return a screenshot.
     """
-    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)
+        logger.info(f"API request received. Extension: {request.extension_percentage}%")
+        # Run the blocking Selenium code in a separate thread (FastAPI handles this)
+        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としてはエラーの方が親切かもしれない
+             logger.error("Screenshot generation failed, returning 1x1 image.")
+             # Optionally return a proper error response instead of 1x1 image
+             # raise HTTPException(status_code=500, detail="Failed to generate screenshot")
 
-        # PIL画像をPNGバイトデータに変換
+
+        # Convert PIL Image to PNG bytes
         img_byte_arr = BytesIO()
         pil_image.save(img_byte_arr, format='PNG')
-        img_byte_arr.seek(0) # ポインタを先頭に戻す
+        img_byte_arr.seek(0)  # Go to the start of the BytesIO buffer
 
         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}")
+        logger.error(f"API Error: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")
 
-# --- Gradioインターフェース定義 ---
-gradio_interface = gr.Interface(
-    fn=render_fullpage_screenshot, # コアロジック関数を再利用
+# --- Gradio Interface Definition ---
+# Note: We reuse the same core function 'render_fullpage_screenshot'
+iface = gr.Interface(
+    fn=render_fullpage_screenshot,
     inputs=[
         gr.Textbox(lines=15, label="HTMLコード入力"),
-        gr.Slider(minimum=0, maximum=30, step=1.0, value=8, label="上下高さ拡張率（%）") # 最大値を少し増やす
+        gr.Slider(minimum=0, maximum=20, step=1.0, value=8, label="上下高さ拡張率（%）")
     ],
     outputs=gr.Image(type="pil", label="ページ全体のスクリーンショット"),
-    title="Full Page Screenshot (高さ拡張調整可能)",
-    description="HTMLをヘッドレスブラウザでレンダリングし、ページ全体を1枚の画像として取得します。上下のみユーザー指定の余裕(%)を追加します。APIは /screenshot (POST) で利用可能です。",
-    allow_flagging="never" # Hugging Face Spacesでフラグ機能を無効化
+    title="Full Page Screenshot (高さ拡張調整可能)", # APIを削除
+    description="HTMLをヘッドレスブラウザでレンダリングし、ページ全体を1枚の画像として取得します。上下のみユーザー指定の余裕(%)を追加します。APIエンドポイントは /api/screenshot で利用可能です。",
+    allow_flagging="never" # Hugging Face Spacesでのフラグ付けを無効化 (任意)
 )
 
-# --- 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)
+# --- Mount Gradio App onto FastAPI ---
+# Mount the Gradio interface at the root path "/"
+app = gr.mount_gradio_app(app, iface, path="/")
+
+# --- Run with Uvicorn (for local testing) ---
+# This part is mainly for local development.
+# When deploying on Hugging Face Spaces, the Spaces runtime handles launching the app.
+if __name__ == "__main__":
+    import uvicorn
+    logger.info("Starting Uvicorn server for local development...")
+    # Host '0.0.0.0' makes it accessible on the network
+    # Port 7860 is a common default for Gradio, but 8000 is common for FastAPI
+    uvicorn.run(app, host="0.0.0.0", port=7860)
 
 # --- END OF FILE app.py ---