ScreenShot / app.py
IotaCluster's picture
Update app.py
175a26c verified
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):
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.set_window_size(1280, 1280) # Set window size to 1920x1080
wd.get(url)
time.sleep(5)
screenshot = wd.get_screenshot_as_png()
if len(screenshot) < 1000:
print("⚠ Screenshot too small, likely error.")
return None
image = Image.open(BytesIO(screenshot))
# Center crop
width, height = image.size
crop_size_side = 1280
crop_size_top = 1280
left = (width - crop_size_side) // 2
top = (height - crop_size_top) // 2
right = left + 1000
bottom = top - 500
cropped_image = image.crop((left, top, right, bottom))
return image
except WebDriverException as e:
print("WebDriverException:", e)
return None
finally:
if wd:
wd.quit()
iface = gr.Interface(
fn=take_screenshot,
inputs=gr.Textbox(label="Website URL", value="https://github.com"),
outputs=gr.Image(type="pil", label="Cropped Screenshot"),
title="Website Screenshot",
description="Takes a 1920x1080 screenshot of the website, center-cropped to 900x900."
)
iface.launch()