Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| 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)) | |
| 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.implicitly_wait(15) | |
| wd.get(url) | |
| 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() | |
| iface = gr.Interface( | |
| fn=take_screenshot, | |
| 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() |