image.py

from skimage import color, img_as_float
from PIL import Image
import numpy as np
import datetime
import pyautogui
from mss import mss

counter = 0
redCarSide = 0
blueCarSide = 1


class Shape:
    def __init__(self, shape, color, side, x, y):
        self.shape = shape
        self.color = color
        self.side = side
        self.x = x
        self.y = y

    def __str__(self):
        return str(self.shape) + ',' + str(self.color) + ',' + str(
            self.side) + ',' + str(self.x) + ',' + str(self.y)


def image_processing():
    global counter
    shapes = []
    redShape = Shape(2, 1, 0, 0, 0)
    blueShape = Shape(2, 1, 0, 0, 0)
    im = screen()
    image = img_as_float(im)
    time1 = datetime.datetime.now()

    distance_red = color.rgb2gray(
        1 - np.abs(image - (0.93359375, 0.20703125, 0.38671875)))
    distance_blue = color.rgb2gray(1 - np.abs(image -
                                              (0.08984375, 0.6875, 0.78125)))
    black_mask_red = color.rgb2gray(distance_red) < 0.95
    black_mask_blue = color.rgb2gray(distance_blue) < 0.99
    white_mask_red = color.rgb2gray(distance_red) > 0.95
    white_mask_blue = color.rgb2gray(distance_blue) > 0.99
    distance_red[black_mask_red] = 0
    distance_blue[black_mask_blue] = 0
    distance_red[white_mask_red] = 1
    distance_blue[white_mask_blue] = 1

    time2 = datetime.datetime.now()
    y = 0
    while y < len(distance_blue):
        for x in [151, 152, 153, 155, 156, 157, 211, 212, 213, 214, 215, 216]:
            if distance_blue[y][x] == 1:
                if 155 <= x <= 157:
                    shapes.append(Shape(0, 1, 0, x, y))
                    blueShape = Shape(0, 1, 0, x, y)
                elif 151 <= x <= 153:
                    shapes.append(Shape(1, 1, 0, x, y))
                    blueShape = Shape(1, 1, 0, x, y)
                elif 214 <= x <= 216:
                    shapes.append(Shape(0, 1, 1, x, y))
                    blueShape = Shape(0, 1, 1, x, y)
                elif 211 <= x <= 213:
                    shapes.append(Shape(1, 1, 1, x, y))
                    blueShape = Shape(1, 1, 1, x, y)
                y += 30
        x = 0
        y += 1

    y = 0
    x = 0
    while y < len(distance_red):
        for x in [30, 31, 32, 34, 35, 36, 94, 95, 96, 90, 91, 92]:
            if distance_red[y][x] == 1:
                if 34 <= x <= 36:
                    shapes.append(Shape(0, 0, 0, x, y))
                    redShape = Shape(0, 0, 0, x, y)
                elif 30 <= x <= 32:
                    shapes.append(Shape(1, 0, 0, x, y))
                    redShape = Shape(1, 0, 0, x, y)
                elif 95 <= x <= 96:
                    shapes.append(Shape(0, 0, 1, x, y))
                    redShape = Shape(0, 0, 1, x, y)
                elif 90 <= x <= 94:
                    shapes.append(Shape(1, 0, 1, x, y))
                    redShape = Shape(1, 0, 1, x, y)
                x = 0
                y += 30
        x = 0
        y += 1

    time3 = datetime.datetime.now()
    print('++++++++++++++++++++' + str(time2 - time1))
    print('--------------------' + str(time3 - time2))

    return redShape, blueShape


def move(red, blue):
    global redCarSide
    global blueCarSide
    global counter
    print('c::::' + str(counter))
    print('redside:::' + str(redCarSide) + '------' + 'red:::' + str(red))
    print('blueside:::' + str(blueCarSide) + '------' + 'blue:::' + str(blue))
    if red:
        redCarSide = 1 if redCarSide == 0 else 0
        pyautogui.moveTo(860, 600)
        pyautogui.mouseDown(button='left')
        pyautogui.mouseUp(button='left')
    if blue:
        blueCarSide = 1 if blueCarSide == 0 else 0
        pyautogui.moveTo(1080, 600)
        pyautogui.mouseDown(button='left')
        pyautogui.mouseUp(button='left')


def screen():
    """ Capture a part of the screen. """

    try:
        with mss() as sct:
            time1 = datetime.datetime.now()

            # The screen part to capture
            mon = {'top': 380, 'left': 840, 'width': 240, 'height': 240}

            # Here we actually capture the screen part
            pixels = sct.get_pixels(mon)

            # Useful informations are available via sct.XXX attributes
            #  see https://python-mss.readthedocs.io/en/dev/api.html#attributes
            # Using these attributes, you can directly create the Image:
            im = Image.frombytes('RGB', (sct.width, sct.height), sct.image)
            # OR, as you already know the picture size:
            im = Image.frombytes('RGB', (240, 240), pixels)

            time2 = datetime.datetime.now()
            print('===========' + str(time2 - time1))

            # And return the Image
            return im
    except Exception as ex:
        print(ex)
        return None


if __name__ == '__main__':
    pyautogui.moveTo(960, 570)
    pyautogui.mouseDown(button='left')
    pyautogui.mouseUp(button='left')

    while True:
        try:
            counter += 1
            redShouldMove = 0
            blueShouldMove = 0
            cal1 = datetime.datetime.now()
            redShape, blueShape = image_processing()
            print(str(counter) + '::::::::::')
            print(redShape)
            print(blueShape)
            if redShape.shape != 2:
                if redCarSide == redShape.side and redShape.shape == 1:
                    redShouldMove = 1
                if redCarSide != redShape.side and redShape.shape == 0:
                    redShouldMove = 1
            if blueShape.shape != 2:
                if blueCarSide == blueShape.side and blueShape.shape == 1:
                    blueShouldMove = 1
                if blueCarSide != blueShape.side and blueShape.shape == 0:
                    blueShouldMove = 1
            cal2 = datetime.datetime.now()
            print('TTTTTTTTTTTTT::::' + str(cal2 - cal1))
            move(redShouldMove, blueShouldMove)
        except KeyboardInterrupt:
            break