File size: 1,891 Bytes
0bba599
 
 
f5611b1
0bba599
f5611b1
 
 
5ca8483
 
0bba599
 
 
 
f5611b1
58b06a1
f5611b1
 
 
0bba599
 
 
 
 
 
f5611b1
0bba599
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f5611b1
 
 
0bba599
 
 
 
f5611b1
5ca8483
 
 
 
57c646c
 
 
 
5ca8483
 
 
f5611b1
 
5ca8483
0bba599
 
 
 
f5611b1
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import time
from io import BytesIO

import gradio as gr
from PIL import Image
from selenium import webdriver
from selenium.common.exceptions import WebDriverException

import main_script

driver = None


def get_chrome_options():
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')

    return options


def set_driver():
    options = get_chrome_options()

    try:
        web_driver = webdriver.Chrome(options=options)
        web_driver.set_window_size(1080, 720)  # Adjust the window size here
    except WebDriverException as e:
        return Image.new('RGB', (1, 1))

    return web_driver


def take_screenshot(url):
    global driver

    try:
        driver = set_driver()  # Adjust the window size here
        driver.get(url)
        driver.implicitly_wait(10)

        images = []
        # Take multiple screenshots (e.g. scroll and capture 3 times)
        for i in range(3):
            png = driver.get_screenshot_as_png()
            image = Image.open(BytesIO(png))
            images.append(image)
            driver.execute_script("window.scrollBy(0, window.innerHeight);")
            time.sleep(1)


    except WebDriverException as e:
        return Image.new('RGB', (1, 1))
    finally:
        if driver:
            driver.quit()

    return images

def call_main_script():
    main_script.main()

def main(url):
    if url== "amit":
        return call_main_script()
    else:
        return take_screenshot(url)




iface = gr.Interface(
    fn=main,
    inputs=gr.Textbox(label="Website URL", value="https://www.google.com/"),
    outputs=gr.Gallery(label="Screenshots", columns=3, height="auto"),
    title="Website Screenshots",
    description="Take multiple screenshots of a website using Selenium in a Gradio Space!"
)

iface.launch()