Selenium and Geckodriver browser screenshot

screenshot.py

import threading
import traceback
import subprocess
import os
import datetime

from pyvirtualdisplay.smartdisplay import SmartDisplay
from selenium import webdriver
from PIL import Image, ImageDraw, ImageFont


def get_screen_resolution():
    cmd = 'xrandr | grep "\*" | cut -d" " -f4'
    output = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE).communicate()[0].decode().strip()
    resolutions = output.split('\n')
    resolution = list(map(lambda r: r.split('x'), resolutions))[0]
    return {'width': int(resolution[0]), 'height': int(resolution[1])}


def screenshot(team_nr, url, filepath, find_elem_by_id=None, scroll_percent=0, verbose=True):
    with screen_lock:
        try:
            resolution = get_screen_resolution()
            with SmartDisplay(visible=0, size=(resolution['width'], resolution['height'])) as display:
                # download geckodriver from: https://github.com/mozilla/geckodriver/releases
                browser = webdriver.Firefox(executable_path=os.getcwd() + '/geckodriver')
                browser.get(url)

                if find_elem_by_id:
                    browser.find_element_by_id(find_elem_by_id)

                if scroll_percent > 0:
                    browser.execute_script("window.scrollTo(0, document.body.scrollHeight*{0:.2f});".format(scroll_percent))

                img = display.waitgrab(1)
                img.save(filepath)
                browser.quit()

                timestamp = datetime.datetime.utcnow().replace(microsecond=0).isoformat() + "Z bt{:02}o{:03}".format(team_nr,  0)

                im = Image.open(filepath)
                draw = ImageDraw.Draw(im)
                draw.text([im.size[0] - (len(timestamp)+1)*10, 10], timestamp)
                del draw

                #write image
                im.save(filepath, 'png')

        except Exception:
            traceback.print_exc()

screen_lock = threading.Lock()

filename = os.path.join('.', 'test.png')
screenshot(79, 'http://example.com', filename)