screen.py

import gradio as gr
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from PIL import Image
from io import BytesIO
import os
import subprocess
proxy_extension_path = os.path.abspath("proxy_extension")
def take_screenshot(url):
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')
    options.add_argument(f'--load-extension={proxy_extension_path}')
    #subprocess.Popen(["x11vnc", "-forever", "-display", ":0", "-rfbauth", "~/.vnc/passwd"])
    try:
        wd = webdriver.Chrome(options=options)
        wd.set_window_size(1080, 720)  # Adjust the window size here
        wd.get(url)
        wd.implicitly_wait(10)
        screenshot = wd.get_screenshot_as_png()
    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.inputs.Textbox(label="Website URL", default="https://kargaranamir.github.io"),
    outputs=gr.Image(type="pil", height=360, width=540),  # Adjust the image size here
    title="Website Screenshot",
    description="Take a screenshot of a website.",
)

iface.launch()