Spaces:
Running
Running
File size: 1,367 Bytes
602cdfd c3b7bd1 602cdfd c3b7bd1 c602fc5 c3b7bd1 602cdfd c3b7bd1 602cdfd c3b7bd1 602cdfd |
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 |
from flask import Flask, request, send_file
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from PIL import Image
from io import BytesIO
app = Flask(__name__)
def take_screenshot(url):
options = webdriver.ChromeOptions()
options.add_argument('--headless')
options.add_argument('--no-sandbox')
options.add_argument('--disable-dev-shm-usage')
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:
# If there's an error, return a blank image as a placeholder
return Image.new('RGB', (1, 1))
finally:
if wd:
wd.quit()
return Image.open(BytesIO(screenshot))
@app.route('/screenshot', methods=['GET'])
def screenshot():
url = request.args.get('url')
if not url:
return "URL parameter is required.", 400
# Take the screenshot of the provided URL
image = take_screenshot(url)
# Save the screenshot to an in-memory file and return as response
img_io = BytesIO()
image.save(img_io, 'PNG')
img_io.seek(0)
return send_file(img_io, mimetype='image/png')
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)
|