Volovikov_CV/trex/main.py

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


cv2.namedWindow('View', cv2.WINDOW_NORMAL)
DEFAULT_MONITOR = 3
scheduler = sched.scheduler(time, sleep)


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)
    binaryScreen = cv2.dilate(binaryScreen, None, iterations=1)
    return binaryScreen


def binary_to_RGB(image):
    rgb_image = np.zeros((*image.shape, 3))
    rgb_image[image == 1] = (255, 255, 255)
    return rgb_image


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()+15*mh, 'width':(y.max()-y.min())+4*mw, 'height':(x.max()-x.min())+10*mh}
    return playing_bbox, mw, mh


def to_bbox(bin_image):
    labels, labeledScreen = cv2.connectedComponents(bin_image.astype('uint8'))
    for label in range(1, labels):
        x,y = np.where(labeledScreen==label)
        bin_image[x.min():x.max()+1, y.min():y.max()+1] = 1
    return bin_image





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

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

    playing_bbox, mw, mh = get_playing_bbox(sct)
    dino_pos = None
    awaitable = None

    binScreen = take_binaryScreenshot(sct, playing_bbox)
    binScreen = to_bbox(binScreen)
    labels, labeledScreen = cv2.connectedComponents(binScreen.astype('uint8'))
    rgbScreen = binary_to_RGB(binScreen)
    dino = min(range(1,labels), key=lambda x: np.where(labeledScreen==x)[1].min())
    dino_pos = np.where(labeledScreen == dino)[1].max()
    rgbScreen[labeledScreen==dino] = (0, 255, 0)


    while cv2.waitKey(1) != ord('q'):
        binScreen = take_binaryScreenshot(sct, playing_bbox)
        binScreen = to_bbox(binScreen)
        labels, labeledScreen = cv2.connectedComponents(binScreen.astype('uint8'))
        rgbScreen = binary_to_RGB(binScreen)
        for label in range(1, labels):
            pos = np.where(labeledScreen == label)
            if pos[1].max() > dino_pos and (pos[1].max()-pos[1].min()) > mw:
                rgbScreen[labeledScreen == label] = (0, 0, 255)
            else: rgbScreen[labeledScreen == label] = (0, 255, 0)
        cv2.imshow('View', rgbScreen)
        cv2.waitKey(1)
    
    cv2.destroyAllWindows()