Volovikov_CV/trex/main.py

import webbrowser
from mss import mss
import numpy as np
from time import sleep
import pyautogui as gui
import cv2


cv2.namedWindow('View', cv2.WINDOW_NORMAL)
DEFAULT_MONITOR = 3


def take_screenshot(sct, monitor=DEFAULT_MONITOR):
    monitor = sct.monitors[monitor]
    image = np.array(sct.grab(monitor))
    return image


def take_binaryScreenshot(sct, monitor: dict):
    image = np.array(sct.grab(monitor))
    binaryScreen = np.zeros(image.shape[:-1])
    binaryScreen[np.all(image == (83, 83, 83, 0), axis=2)] = 1
    binaryScreen = cv2.erode(binaryScreen, None, iterations=1)
    return binaryScreen


def open_dino(sct):
    prevScreen = take_screenshot(sct)
    webbrowser.open('https://chromedino.com/')
    sleep(1)
    Screen = take_screenshot(sct)
    while np.any(Screen != prevScreen):
        cv2.imshow('View', Screen)
        cv2.waitKey(1)
        sleep(1)
        prevScreen = Screen.copy()
        Screen = take_screenshot(sct)
    print('loaded')


def short_jump():
    with gui.hold('space'):
        sleep(0.05)


def long_jump():
    with gui.hold('space'):
        sleep(0.2)


def get_playing_bbox(sct):
    Screen = take_screenshot(sct)
    binaryScreen = np.zeros(Screen.shape[:-1])
    binaryScreen[np.all(Screen == (83, 83, 83, 0), axis=2)] = 1
    short_jump()
    sleep(1)

    binScreen = take_binaryScreenshot(sct, sct.monitors[DEFAULT_MONITOR])
    x, y = np.where(binScreen == 1)
    mw = int((y.max()-y.min()) / 100)
    mh = int((x.max()-x.min()) / 100)
    playing_bbox = {'left':y.min()-2*mw, 'top':x.min()-5*mh, 'width':(y.max()-y.min())+4*mw, 'height':(x.max()-x.min())+10*mh}
    return playing_bbox





if __name__ == '__main__':
    sct = mss()

    open_dino(sct)
    gui.hotkey('fn', 'f11')

    playing_bbox = get_playing_bbox(sct)

    while cv2.waitKey(1) != ord('q'):
        binScreen = take_binaryScreenshot(sct, playing_bbox)
        cv2.imshow('View', binScreen)
        cv2.waitKey(1)
    
    cv2.destroyAllWindows()