File size: 2,088 Bytes
c3b7bd1
 
 
 
 
63c3636
 
 
 
 
 
 
 
 
 
 
 
c3b7bd1
 
 
 
 
63c3636
 
c3b7bd1
 
63c3636
c3b7bd1
63c3636
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c3b7bd1
 
63c3636
c3b7bd1
 
 
 
 
 
63c3636
 
 
 
 
 
 
 
 
c3b7bd1
dea635c
c3b7bd1
 
63c3636
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
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()