Update app.py

app.py

@@ -8,18 +8,16 @@ from PIL import Image
 import time
 import base64
 
-
-def take_screenshot_from_html(html_code, wait_time=3):
+def take_screenshot_from_html(html_code, height_adjustment_percent=3):
     """
     HTMLコードからフルページスクリーンショットを取得する関数。
-    キャプチャ前に高さを3%拡張し、スクロールバーの有無に応じて調整。
 
     Args:
         html_code (str): スクリーンショットを取得するHTMLコード。
-        wait_time (int): ページ読み込みの待機時間（秒）。
+        height_adjustment_percent (float): 高さの調整率（パーセント）。
 
     Returns:
-        tuple: (PIL.Image.Image, str) スクリーンショット画像とメッセージ。
+        tuple: (PIL.Image.Image, str) スクリーンショット画像とメッセージのタプル。
                エラーの場合は (None, str) を返す。
     """
     options = Options()
@@ -29,73 +27,65 @@ def take_screenshot_from_html(html_code, wait_time=3):
     options.add_argument("--disable-gpu")
     options.add_argument("--window-position=-2400,-2400")  # Windowsでのバグ回避
 
-    driver = None  # driverをtryの外で初期化
-
     try:
         driver = webdriver.Chrome(options=options)
         html_base64 = base64.b64encode(html_code.encode('utf-8')).decode('utf-8')
         driver.get(f"data:text/html;base64,{html_base64}")
 
-        time.sleep(wait_time)
+        time.sleep(3)  # 固定の待機時間 (3秒)
 
-        # スクロールバーの有無をチェック
-        has_scrollbar = driver.execute_script(
-            "return document.documentElement.scrollHeight > document.documentElement.clientHeight;"
-        )
-
-        # ページサイズを取得
+        # フルページサイズを取得 (JavaScript)
         width = driver.execute_script(
             "return Math.max(document.body.scrollWidth, document.body.offsetWidth, "
             "document.documentElement.clientWidth, document.documentElement.scrollWidth, "
             "document.documentElement.offsetWidth);"
         )
-        height = driver.execute_script(
+        original_height = driver.execute_script(
             "return Math.max(document.body.scrollHeight, document.body.offsetHeight, "
             "document.documentElement.clientHeight, document.documentElement.scrollHeight, "
             "document.documentElement.offsetHeight);"
         )
 
-        # 高さを3%拡張（スクロールバーがあればさらに調整）
-        extra_height = int(height * 0.03)
-        if has_scrollbar:
-            final_height = min(height + extra_height + 20, 5000)  # スクロールバー分を考慮、最大5000px
-        else:
-             final_height = min(height + extra_height, 5000)  # スクロールバーがなければ3%のみ、最大5000px
+        # 高さを調整
+        adjusted_height = int(original_height * (1 + height_adjustment_percent / 100))
 
-        driver.set_window_size(min(width + 100, 1920), final_height)  # 最大幅1920
+        # ウィンドウサイズを設定（最大サイズ制限付き）
+        driver.set_window_size(min(width + 100, 1920), min(adjusted_height + 100, 5000))
         time.sleep(1)
 
+        # スクリーンショットを一時ファイルに保存
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp:
             temp_filename = temp.name
         driver.save_screenshot(temp_filename)
+
+        driver.quit()
         image = Image.open(temp_filename)
         os.unlink(temp_filename)
         return image, "スクリーンショットを取得しました。"
 
     except (WebDriverException, TimeoutException) as e:
+        driver.quit()
         return None, f"エラー: {str(e)}"
     except Exception as e:
+        driver.quit()
         return None, f"予期しないエラー: {str(e)}"
-    finally:
-        if driver:  # driverが存在する場合のみquit()を実行
-            driver.quit()
 
 
-# Gradioインターフェース (HTML専用)
+# Gradioインターフェース
 with gr.Blocks(title="HTML Screenshot Tool") as demo:
     gr.Markdown("# HTML スクリーンショットツール")
     gr.Markdown("HTMLコードを入力して、レンダリングされたページのフルスクリーンショットを取得します。")
 
     html_input = gr.Textbox(label="HTMLコード", placeholder="<p>Hello, world!</p>", lines=5)
-    wait_time_slider = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="待機時間 (秒)")
-    screenshot_button = gr.Button("スクリーンショットを取得")
-    image_output = gr.Image(label="スクリーンショット")
-    text_output = gr.Textbox(label="結果")
+    height_adjustment_slider = gr.Slider(minimum=0, maximum=10, value=3, step=0.1, label="高さ調整率 (%)")
+    html_button = gr.Button("スクリーンショットを取得")
+    html_image_out = gr.Image(label="スクリーンショット")
+    html_text_out = gr.Textbox(label="結果")
 
-    screenshot_button.click(
+    html_button.click(
         take_screenshot_from_html,
-        inputs=[html_input, wait_time_slider],
-        outputs=[image_output, text_output]
+        inputs=[html_input, height_adjustment_slider],
+        outputs=[html_image_out, html_text_out]
     )
 
 demo.launch()