TangoGame/screenreader.py at main · viccollado/TangoGame · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import numpy as np
import pyautogui
import torch
import torchvision.transforms as transforms

from model import Model

def read_screen_info():
    # Load trained model
    device = torch.device("cpu")
    model = Model().to(device)
    model.load_state_dict(torch.load("cnn.pth", map_location=device))
    model.eval()
    # Define the same transform as used in training
    transform = transforms.ToTensor()


    xtop, ytop = 754, 271
    xbottom, ybottom = 1151, 668

    board = pyautogui.screenshot(region=(xtop, ytop, xbottom - xtop, ybottom - ytop))
    # board.save("screenshot_base3.png")

    xdiff = xbottom - xtop
    ydiff = ybottom - ytop

    xbox_size = xdiff // 6
    ybox_size = ydiff // 6

    # Vertical middle points
    middle_x = [int(xbox_size + i * xbox_size) for i in range(6)]
    middle_y = [int(ybox_size // 2 + i * ybox_size) for i in range(6)]

    # Horizontal middle points
    middle_x2 = [int(xbox_size // 2 + i * xbox_size) for i in range(6)]
    middle_y2 = [int(ybox_size + i * ybox_size) for i in range(6)]

    # 2 arrays the 0,1,2 sobre la info de los ejes.
    edges_array_vertical = np.zeros(shape=(6, 5))  # Matriz 6x5
    edges_array_horizontal = np.zeros(shape=(5, 6))  # Matriz 5x6

    # Predict the values of the possible cross and equals
    vertical_imgs = []
    vertical_indices = []
    for i in range(5):
        for j in range(6):
            minibox = board.crop(
                (middle_x[i] - 10, middle_y[j] - 10, middle_x[i] + 10, middle_y[j] + 10)
            ).convert("RGB")
            vertical_imgs.append(transform(minibox))
            vertical_indices.append((j, i))
            # minibox.save(f"miniboxes/vertical/vbox2_{j}{i}.png")

    vertical_batch = torch.stack(vertical_imgs).to(device)
    with torch.no_grad():
        outputs = model(vertical_batch)
        preds = torch.argmax(outputs, dim=1).cpu().numpy()

    for idx, pred in zip(vertical_indices, preds):
        edges_array_vertical[idx] = pred

    # Batch prediction for horizontal edges
    horizontal_imgs = []
    horizontal_indices = []
    for i in range(6):
        for j in range(5):
            minibox = board.crop(
                (middle_x2[i] - 10, middle_y2[j] - 10, middle_x2[i] + 10, middle_y2[j] + 10)
            ).convert("RGB")
            horizontal_imgs.append(transform(minibox))
            horizontal_indices.append((j, i))
            # minibox.save(f"miniboxes/horizontal/hbox2_{j}{i}.png")

    horizontal_batch = torch.stack(horizontal_imgs).to(device)
    with torch.no_grad():
        outputs = model(horizontal_batch)
        preds = torch.argmax(outputs, dim=1).cpu().numpy()

    for idx, pred in zip(horizontal_indices, preds):
        edges_array_horizontal[idx] = pred

    # Matrices con los valores de los edges
    # print(f"Vertical array edges: \n {edges_array_vertical}\n")
    # print(f"Horizontal array edges: \n {edges_array_horizontal}\n")

    """
    Ejemplo de array para la "screenshot_base.png"

    0: nada
    1: cross
    2: equal

    edges_array_vertical =[
        [0, 0, 2, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0],
    ]

    edges_array_horizontal =[
        [0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 2],
        [0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0]
    ]
    """


    # Check if there are suns or moons
    middle_x3 = [int(xbox_size // 2 + i * xbox_size) for i in range(6)]
    middle_y3 = [int(ybox_size // 2 + i * ybox_size) for i in range(6)]

    sun_indices = []
    moon_indices = []
    # In RGB
    SUN_COLOR = np.asarray([255,179,30])
    MOON_COLOR = np.asarray([130,175,237])

    for i in range(6):
        for j in range(6):
            minibox = board.crop(
                (middle_x3[i]+5, middle_y3[j]+5, middle_x3[i]+6, middle_y3[j]+6)
            ).convert("RGB")
            # minibox.save(f"miniboxes/moonsuns/box_{j}{i}.png")
            sun_distance = np.sqrt(np.sum((SUN_COLOR - np.asarray(minibox))**2))
            moon_distance = np.sqrt(np.sum((MOON_COLOR - np.asarray(minibox))**2))
            if sun_distance <= 1:
                sun_indices.append((j,i))
            if moon_distance <= 1:
                moon_indices.append((j,i))

    # print(sun_indices)
    # print(moon_indices)

    return edges_array_vertical, edges_array_horizontal, sun_indices, moon_indices, middle_x3, middle_y3, xtop, ytop