Update app.py

app.py

@@ -3,33 +3,68 @@ from selenium import webdriver
 from selenium.common.exceptions import WebDriverException
 from PIL import Image
 from io import BytesIO
+import time
+
+def take_screenshot(url, device):
+    
+    presets = {
+        "Mobile (375x667)": (375, 667),
+        "Tablet (768x1024)": (768, 1024),
+        "PC (1366x768)": (1366, 768),
+        "Full Screenshot": (1920, 1080)
+    }
+    
+    width, height = presets.get(device, (1366, 768))
 
-def take_screenshot(url):
     options = webdriver.ChromeOptions()
     options.add_argument('--headless')
     options.add_argument('--no-sandbox')
     options.add_argument('--disable-dev-shm-usage')
-
+    
+    wd = None
     try:
         wd = webdriver.Chrome(options=options)
-        wd.set_window_size(1345, 846)  # Adjust the window size here
+        wd.implicitly_wait(15)
         wd.get(url)
-        wd.implicitly_wait(30)
-        screenshot = wd.get_screenshot_as_png()
+        
+        if device == "Full Screenshot":
+            
+            required_width = wd.execute_script("return document.body.parentNode.scrollWidth")
+            required_height = wd.execute_script("return document.body.parentNode.scrollHeight")
+
+            max_height = 10000 
+            
+            wd.set_window_size(width, min(required_height, max_height))
+            time.sleep(1) 
+            screenshot = wd.get_screenshot_as_png()
+            
+        else:
+            wd.set_window_size(width, height)
+            time.sleep(1) 
+            screenshot = wd.get_screenshot_as_png()
+            
+        return Image.open(BytesIO(screenshot))
+        
     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", value="https://kargaranamir.github.io"),
-    outputs=gr.Image(type="pil", label="Screenshot", height=360, width=540),  # Adjust the image size here
+    inputs=[
+        gr.Textbox(label="Website URL", value="https://asitha.top/"),
+        gr.Dropdown(
+            choices=["Mobile (375x667)", "Tablet (768x1024)", "PC (1366x768)", "Full Screenshot"],
+            label="Device",
+            value="PC (1366x768)"
+        )
+    ],
+    outputs=gr.Image(type="pil", label="Screenshot", height=360, width=540),
     title="Website Screenshot",
     description="Take a screenshot of a website using selenium in a gradio space!"
 )
 
-iface.launch()
+iface.launch()