heboya8/selenium

1import gradio as gr
2from selenium import webdriver
3from selenium.common.exceptions import WebDriverException
4from PIL import Image
5from io import BytesIO
6
7def take_screenshot(url):
8    options = webdriver.ChromeOptions()
9    options.add_argument('--headless')
10    options.add_argument('--no-sandbox')
11    options.add_argument('--disable-dev-shm-usage')
12
13    try:
14        wd = webdriver.Chrome(options=options)
15        wd.set_window_size(1080, 720)  # Adjust the window size here
16        wd.get(url)
17        wd.implicitly_wait(10)
18        screenshot = wd.get_screenshot_as_png()
19    except WebDriverException as e:
20        return Image.new('RGB', (1, 1))
21    finally:
22        if wd:
23            wd.quit()
24
25    return Image.open(BytesIO(screenshot))
26
27iface = gr.Interface(
28    fn=take_screenshot,
29    inputs=gr.Textbox(label="Website URL", value="https://kargaranamir.github.io"),
30    outputs=gr.Image(type="pil", label="Screenshot", height=360, width=540),  # Adjust the image size here
31    title="Website Screenshot",
32    description="Take a screenshot of a website using selenium in a gradio space!"
33)
34
35iface.launch()
36