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
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__':
    with mss() as sct:

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

        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()
        dino_max_y = dino_pos[0].min()
        dino_min_x = dino_pos[1].min()

        coef = 3
        trigger_time = (int(dino_pos[0].mean()), dino_max_x + 18*pw)
        trigger_jump = (int(dino_pos[0].mean()), dino_max_x + 8*pw)
        trigger_feather_time = (dino_max_y-5*ph, dino_max_x + 23*pw)
        trigger_feather_take = (dino_max_y-5*ph, dino_max_x + 13*pw)
        
        enter_time = None
        enter_time_feather = None
        jump_time = None
        take_time = None

        start = time()
        start2 = time()

        last_screen = None
        speed = 0

        stop_speeding = False
                                
        if DEBUG: cv2.imshow('View', rgbScreen) 
        while True:
            binScreen = take_binaryScreenshot(sct, playing_bbox)                     
            binScreen = to_bbox(binScreen)
            labels, labeledScreen = cv2.connectedComponents(binScreen.astype('uint8'))

            if not enter_time and binScreen[trigger_time] != 0:
                enter_time = time()

            elif enter_time and binScreen[trigger_time] == 0:
                jump_time = (time() - enter_time) / coef
                enter_time = None

            if binScreen[trigger_jump] != 0:
                gui.press('space')
                if jump_time:
                    sleep(jump_time)
                    gui.press('down') 
                    jump_time = None

            if not enter_time_feather and binScreen[trigger_feather_time] != 0:
                enter_time_feather = time()

            elif enter_time_feather and binScreen[trigger_feather_time] == 0:
                take_time = time() - enter_time_feather
                enter_time_feather = None

            if binScreen[trigger_feather_take] != 0:
                print('take')
                print(take_time)
                with gui.hold('down'):
                    sleep(take_time*4)
                take_time = None
            
            if (time() - start)*coef > 25*3.3 and stop_speeding == False:
                if (time() - start) < 10.2: stop_speeding = True
                print(time() - start)
                trigger_jump = (trigger_jump[0], trigger_jump[1]+pw)
                trigger_time = (trigger_time[0], trigger_time[1]+pw)
                trigger_feather_take = (trigger_feather_take[0], trigger_feather_take[1]+int(pw*1.4))
                trigger_feather_time = (trigger_feather_time[0], trigger_feather_time[1]+int(pw*1.4))
                start = time()

            if (time() - start2)*coef > 9*3.3 and stop_speeding == False:
                print(time() - start2)
                coef += 0.2
                start2 = time()

            if DEBUG: 
                rgbScreen = binary_to_RGB(binScreen)
                cv2.circle(rgbScreen, trigger_time[::-1], radius=2, color=(255, 0, 0), thickness=2)
                cv2.circle(rgbScreen, trigger_jump[::-1], radius=2, color=(0, 255, 0), thickness=2)
                cv2.circle(rgbScreen, trigger_feather_time[::-1], radius=2, color=(255, 0, 0), thickness=2)
                cv2.circle(rgbScreen, trigger_feather_take[::-1], radius=2, color=(0, 255, 0), thickness=2)
                cv2.imshow('View', rgbScreen)
                if cv2.waitKey(1) == ord('q'): break

            # print(speed)
            # last_screen = binScreen.copy()
            

        
        cv2.destroyAllWindows()