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')

        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()

        coef = 1
        trigger_second = (int(dino_pos[0].mean()), int(dino_max_x + int(coef*10)*pw))
        trigger_feather = (dino_max_y, int(dino_max_x + int(coef*12)*pw))
        jumping = False
        taking = False
        laing = False

        print(ph, pw)

        start = time()
        start2= time()
                                
        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 jumping and binScreen[trigger_second] != 0 or binScreen[trigger_feather] != 0: 
                print('jumping')
                gui.keyDown('space')
                jumping = True

            elif jumping and binScreen[trigger_second[0], trigger_second[1]+(int(coef-1.2)*pw)] == 0 and binScreen[trigger_feather[0], trigger_feather[1]+(int(coef-1.2)*pw)] == 0:
                print('taking')
                gui.keyUp('space')
                gui.keyDown('down')
                jumping = False
                taking = True
            
            if taking and binScreen[dino_max_y-ph, dino_max_x] != 0:
                print('take off')
                gui.keyUp('down')
                taking = False

            if time() - start2 > 15/coef:
                print(coef)
                print('speeded')
                coef += 0.2
                start2 = time()

            if time() - start > 15:
                print(time() - start)
                trigger_second = (trigger_second[0], trigger_second[1]+pw)
                trigger_feather = (trigger_feather[0], trigger_feather[1]+pw)
                start = time()

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

        
        cv2.destroyAllWindows()