Update app.py

app.py

@@ -1,81 +1,59 @@
 import gradio as gr
+import helium
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
-from selenium.webdriver.common.by import By
-import time
-import os
+from io import BytesIO
+from PIL import Image
 
-def capture_screenshot(html_content, width=1280, height=720):
+def capture_screenshot(url):
     """
-    指定されたHTMLコンテンツのスクリーンショットを撮る関数。
-
-    Args:
-        html_content: スクリーンショットを撮りたいHTMLコンテンツの文字列。
-        width: ブラウザウィンドウの幅 (ピクセル単位)。
-        height: ブラウザウィンドウの高さ (ピクセル単位)。
-
-    Returns:
-        スクリーンショット画像のファイルパス。
-        エラーが発生した場合はNoneを返す。
+    指定された URL のスクリーンショットを撮り、PIL Image オブジェクトとして返す。
     """
     try:
-        # Headless Chrome のオプションを設定
+        # Chrome options for headless browsing
         chrome_options = Options()
-        chrome_options.add_argument("--headless")  # ヘッドレスモードを有効にする
-        chrome_options.add_argument("--no-sandbox") # sandboxを無効にする（rootで実行する場合に必要）
-        chrome_options.add_argument("--disable-dev-shm-usage") # /dev/shm を使わないようにする
-        chrome_options.add_argument(f"--window-size={width},{height}")  # ウィンドウサイズ
-
-        #chromedriverのパスを指定（/usr/bin/chromedriver が標準）
-        chromedriver_path = "/usr/bin/chromedriver"  # Hugging Face Spacesの標準パス
-
-        # Chrome WebDriver を初期化
-        driver = webdriver.Chrome(executable_path=chromedriver_path, options=chrome_options)
-
-        # HTMLコンテンツをファイルに書き込む（一時ファイル）
-        temp_html_path = "temp.html"
-        with open(temp_html_path, "w", encoding="utf-8") as f:
-            f.write(html_content)
-
+        chrome_options.add_argument("--headless")  # Headless mode
+        chrome_options.add_argument("--no-sandbox") # For security reasons on some platforms
+        chrome_options.add_argument("--disable-dev-shm-usage") # Fixes /dev/shm out-of-memory issues
+        chrome_options.add_argument("--window-size=1024,768")  # Set a reasonable window size
 
-        # ブラウザで一時HTMLファイルを開く
-        driver.get("file://" + os.path.abspath(temp_html_path))
-        time.sleep(1)  # 念のため少し待つ（JavaScriptの実行などを考慮）
+        # Start the browser with helium (using the options)
+        driver = helium.start_chrome(options=chrome_options)
+        helium.go_to(url)
 
-        # スクリーンショットを撮る
-        screenshot_path = "screenshot.png"
-        driver.save_screenshot(screenshot_path)
+        # Get screenshot as bytes
+        png_bytes = driver.get_screenshot_as_png()
 
-        # ブラウザを閉じる
-        driver.quit()
+        # Convert bytes to PIL Image
+        image = Image.open(BytesIO(png_bytes))
 
-        # 一時HTMLファイルを削除
-        os.remove(temp_html_path)
-
-        return screenshot_path
+        # Kill the browser
+        helium.kill_browser()
+        return image
 
     except Exception as e:
-        print(f"Error: {e}")
-        return None
-
-
-# Gradio インターフェースの定義
-iface = gr.Interface(
-    fn=capture_screenshot,
-    inputs=[
-        gr.HTML(label="HTML Content"),  # HTML入力用のコンポーネント
-        gr.Slider(minimum=320, maximum=1920, value=1280, label="Width"),  # 幅指定
-        gr.Slider(minimum=240, maximum=1080, value=720, label="Height"), # 高さ指定
-
-    ],
-    outputs=gr.Image(type="filepath", label="Screenshot"), # 画像出力
-    title="HTML Screenshot Capture",
-    description="Enter HTML content and capture a screenshot.",
-    examples=[
-        ["<h1>Hello, World!</h1><p>This is a test.</p>", 800, 600], # サンプル
-        ["<div style='background-color: lightblue; padding: 20px;'><h2>Styled Div</h2></div>", 640, 480], # サンプル
-    ]
-)
-
-if __name__ == "__main__":
-    iface.launch()
+        # Handle errors gracefully.
+        return f"Error: {e}"
+    finally:
+        try:
+            helium.kill_browser()
+        except:
+            pass
+
+
+
+# Gradio Interface
+if __name__ == '__main__':
+    iface = gr.Interface(
+        fn=capture_screenshot,
+        inputs=gr.Textbox(label="Enter URL", placeholder="https://www.example.com"),
+        outputs=gr.Image(type="pil", label="Screenshot"), # Use PIL Image output
+        title="HTML Screenshot Capture",
+        description="Enter a URL to capture a screenshot of the webpage.",
+        examples=[
+            ["https://www.google.com"],
+            ["https://www.wikipedia.org"],
+            ["https://www.github.com"]
+        ],
+    )
+    iface.launch(server_name="0.0.0.0", server_port=7860)