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
from collections import namedtuple


cv2.namedWindow('View', cv2.WINDOW_NORMAL)
DEFAULT_MONITOR = 3
DEBUG = True
scheduler = sched.scheduler(time, sleep)
dilation_kernel = np.array([[0,0,0], [1, 1, 1], [0,0,0]]).astype('uint8')


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 shortest_jump():
    with gui.hold('space'):
        sleep(0.1)
    take_down()


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


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


def take_down():
    with gui.hold('down'):
        sleep(0.04)


def jump(sec):
    print('jump suka jump')
    with gui.hold('space'):
        sleep(sec)
    take_down()


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 = (y.max()-y.min()) // 100
    mh = (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):
    bin_image = cv2.dilate(bin_image, dilation_kernel, iterations=4)
    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, pw, ph = get_playing_bbox(sct)

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

    coef = 1.2
    trigger_second = (int(dino_pos[0].mean()), int(dino_max_x + int(coef*10)*pw))
    trigger_first = (int(dino_pos[0].mean()), int(75*pw))
    jumping = False
    taking = False

    start = time()

    cv2.imshow('View', rgbScreen) 
    while cv2.waitKey(1) != ord('q'):
        binScreen = take_binaryScreenshot(sct, playing_bbox)
        binScreen = to_bbox(binScreen)
        labels, labeledScreen = cv2.connectedComponents(binScreen.astype('uint8'))

        if not taking and not jumping and binScreen[trigger_second] != 0:
            print('jumping')
            gui.keyDown('space')
            jumping = True

        elif jumping and binScreen[trigger_second[0], trigger_second[1]-pw] == 0:
            print('landing')
            gui.keyUp('space')
            gui.keyDown('down')
            jumping = False
            taking = True
        
        elif taking and binScreen[dino_min_y, dino_max_x] != 0:
            print('landed')
            gui.keyUp('down')
            taking = False

        if time() - start > 10:
            print(time() - start)
            coef += 0.1
            trigger_second = (trigger_second[0], trigger_second[1]+pw)
            # trigger_first = (trigger_first[0], trigger_first[0]+int(pw))
            start = time()

        if DEBUG: 
            rgbScreen = binary_to_RGB(binScreen)
            cv2.circle(rgbScreen, trigger_second[::-1], radius=2, color=(0, 255, 0), thickness=2)
            cv2.circle(rgbScreen, trigger_first[::-1], radius=2, color=(0, 0, 255), thickness=2)
            cv2.imshow('View', rgbScreen)
        

    
    cv2.destroyAllWindows()