Update app.py

app.py

@@ -1,8 +1,7 @@
 import gradio as gr
 from selenium import webdriver
-from selenium.common.exceptions import WebDriverException
-from selenium.webdriver.chrome.service import Service
 from webdriver_manager.chrome import ChromeDriverManager
+from selenium.webdriver.chrome.service import Service
 from PIL import Image
 from io import BytesIO
 
@@ -12,39 +11,31 @@ def take_screenshot(url, capture_type):
     options.add_argument('--no-sandbox')
     options.add_argument('--disable-dev-shm-usage')
 
-    wd = None
-    try:
-        service = Service(ChromeDriverManager().install())
-        wd = webdriver.Chrome(service=service, options=options)
-        wd.get(url)
-        wd.implicitly_wait(10)
+    wd = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=options)
 
-        if capture_type == "전체 페이지":
-            # 전체 페이지 높이와 너비를 가져옵니다.
-            total_width = wd.execute_script("return document.body.offsetWidth")
-            total_height = wd.execute_script("return document.body.parentNode.scrollHeight")
-            wd.set_window_size(total_width, total_height)  # 윈도우 크기를 전체 페이지 크기로 조정합니다.
-        else:
-            wd.set_window_size(1024, 768)  # 기본 크기
+    if capture_type == "전체 페이지":
+        # 전체 페이지 캡처 로직 (예시: JavaScript를 사용하여 페이지 높이 측정)
+        wd.get(url)
+        total_height = wd.execute_script("return document.body.parentNode.scrollHeight")
+        wd.set_window_size(1024, total_height)  # 창 크기를 전체 페이지 높이에 맞춤
+    else:
+        wd.set_window_size(1024, 768)  # 기본 크기
 
-        screenshot = wd.get_screenshot_as_png()
-    except WebDriverException as e:
-        return Image.new('RGB', (1, 1))
-    finally:
-        if wd:
-            wd.quit()
+    wd.get(url)
+    screenshot = wd.get_screenshot_as_png()
+    wd.quit()
 
     return Image.open(BytesIO(screenshot))
 
 iface = gr.Interface(
     fn=take_screenshot,
     inputs=[
-        gr.inputs.Textbox(label="Website URL", default="https://korating.com"),
-        gr.inputs.Radio(choices=["기본 크기", "전체 페이지"], label="캡처 유형 선택")
+        gr.Textbox(label="Website URL", default="https://korating.com"),
+        gr.Radio(choices=["기본 크기", "전체 페이지"], label="캡처 유형 선택")
     ],
-    outputs=gr.Image(type="pil", tool="editor"),
+    outputs=gr.Image(type="pil"),
     title="Website Screenshot",
-    description="Take a screenshot of a website. Choose '전체 페이지' to capture the full page.",
+    description="Take a screenshot of a website."
 )
 
 iface.launch()