Update app.py

app.py

@@ -1,42 +1,36 @@
 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 PIL import Image
 from io import BytesIO
 
-def take_screenshot(url, capture_type):
+def take_screenshot(url):
     options = webdriver.ChromeOptions()
     options.add_argument('--headless')
     options.add_argument('--no-sandbox')
     options.add_argument('--disable-dev-shm-usage')
 
-    wd = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=options)
-
-    if capture_type == "전체 페이지":
+    wd = None
+    try:
+        wd = webdriver.Chrome(options=options)
+        wd.set_window_size(1920, 1080)  # 조정된 윈도우 크기
         wd.get(url)
-        wd.set_window_size(1920, 1080)  # 임시로 큰 값을 설정합니다. 실제 높이는 아래에서 계산됩니다.
-        total_height = wd.execute_script("return document.body.scrollHeight")
-        wd.set_window_size(1920, total_height)  # 페이지 전체 높이로 창 크기 조정
-    else:
-        wd.set_window_size(1920, 1080)  # 기본 크기
-
-    wd.get(url)
-    screenshot = wd.get_screenshot_as_png()
-    wd.quit()
+        wd.implicitly_wait(10)
+        screenshot = wd.get_screenshot_as_png()
+    except WebDriverException as e:
+        return Image.new('RGB', (1, 1))
+    finally:
+        if wd:
+            wd.quit()
 
     return Image.open(BytesIO(screenshot))
 
 iface = gr.Interface(
     fn=take_screenshot,
-    inputs=[
-        gr.Textbox(label="Website URL", placeholder="https://korating.com"),
-        gr.Radio(choices=["기본 크기", "전체 페이지"], label="캡처 유형 선택", default="기본 크기")
-    ],
-    outputs=gr.Image(type="pil"),
-    title="Website Screenshot Tool",
-    description="Take a screenshot of a website. Choose between default size or full page."
+    inputs=gr.inputs.Textbox(label="Website URL", default="https://korating.com"),
+    outputs=gr.Image(type="pil", height=540, width=960),  # 조정된 이미지 크기
+    title="Website Screenshot",
+    description="Take a screenshot of a website.",
 )
 
-iface.launch()
+iface.launch()