Update app.py

app.py

@@ -1,10 +1,12 @@
 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):
+def take_screenshot(url, capture_type):
     options = webdriver.ChromeOptions()
     options.add_argument('--headless')
     options.add_argument('--no-sandbox')
@@ -12,10 +14,19 @@ def take_screenshot(url):
 
     wd = None
     try:
-        wd = webdriver.Chrome(options=options)
-        wd.set_window_size(1920, 1080)  # 조정된 윈도우 크기
+        service = Service(ChromeDriverManager().install())
+        wd = webdriver.Chrome(service=service, options=options)
         wd.get(url)
         wd.implicitly_wait(10)
+
+        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)  # 기본 크기
+
         screenshot = wd.get_screenshot_as_png()
     except WebDriverException as e:
         return Image.new('RGB', (1, 1))
@@ -27,10 +38,13 @@ def take_screenshot(url):
 
 iface = gr.Interface(
     fn=take_screenshot,
-    inputs=gr.inputs.Textbox(label="Website URL", default="https://korating.com"),
-    outputs=gr.Image(type="pil", height=540, width=960),  # 조정된 이미지 크기
+    inputs=[
+        gr.inputs.Textbox(label="Website URL", default="https://korating.com"),
+        gr.inputs.Radio(choices=["기본 크기", "전체 페이지"], label="캡처 유형 선택")
+    ],
+    outputs=gr.Image(type="pil", tool="editor"),
     title="Website Screenshot",
-    description="Take a screenshot of a website.",
+    description="Take a screenshot of a website. Choose '전체 페이지' to capture the full page.",
 )
 
 iface.launch()