-
Notifications
You must be signed in to change notification settings - Fork 12
/
board_plot.py
159 lines (123 loc) · 5.71 KB
/
board_plot.py
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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
import cv2
import numpy as np
from time import time
def board_plot(board_config, board_state, is_display=False, save_dir=None):
t0 = time()
# Board config
GRID_SIZE = board_config["GRID_SIZE"]
SQUARE_SIZE = board_config["SQUARE_SIZE"]
LINE_THICKNESS = board_config["LINE_THICKNESS"]
BACKGROUND_COLOR = board_config["BACKGROUND_COLOR"]
LINES_COLOR = board_config["LINES_COLOR"]
SNAKE1_COLOR = board_config["SNAKE1_COLOR"]
SNAKE2_COLOR = board_config["SNAKE2_COLOR"]
FOOD_COLOR = board_config["FOOD_COLOR"]
# Board state
turn = board_state["turn"]
snake1_body = board_state["snake1"]["body"]
snake1_dir = board_state["snake1"]["dir"]
snake2_body = board_state["snake2"]["body"]
snake2_dir = board_state["snake2"]["dir"]
food = board_state["food"]
grid_thickness = GRID_SIZE * SQUARE_SIZE + (GRID_SIZE + 1) * LINE_THICKNESS
# Create a white background image
image = np.ones((grid_thickness, grid_thickness, 3), dtype=np.uint8) * BACKGROUND_COLOR
image = image.astype(np.uint8)
# Draw the grid lines
cv2.line(image, (0, 0), (0, grid_thickness), LINES_COLOR, thickness=LINE_THICKNESS)
cv2.line(image, (0, 0), (grid_thickness, 0), LINES_COLOR, thickness=LINE_THICKNESS)
for i in range(1, GRID_SIZE + 1):
# Vertical lines
v_start_point = (i * (SQUARE_SIZE + LINE_THICKNESS), 0)
v_end_point = (i * (SQUARE_SIZE + LINE_THICKNESS), grid_thickness)
cv2.line(image, v_start_point, v_end_point, LINES_COLOR, thickness=LINE_THICKNESS)
# Horizontal lines
h_start_point = (0, i * (SQUARE_SIZE + LINE_THICKNESS))
h_end_point = (grid_thickness, i * (SQUARE_SIZE + LINE_THICKNESS))
cv2.line(image, h_start_point, h_end_point, LINES_COLOR, thickness=LINE_THICKNESS)
# Function to draw every snake and food square position in the grid
def draw_pos(target_pos, color, is_head=False, dir=None):
# Getting the square's coordinates
square_top_left = (target_pos[0] * (SQUARE_SIZE + LINE_THICKNESS) + LINE_THICKNESS, target_pos[1] * (SQUARE_SIZE + LINE_THICKNESS) + LINE_THICKNESS)
square_bottom_right = (square_top_left[0] + SQUARE_SIZE - LINE_THICKNESS, square_top_left[1] + SQUARE_SIZE - LINE_THICKNESS)
cv2.rectangle(image, square_top_left, square_bottom_right, color, -1)
# Drawing the eyes in the head
if is_head:
radius = SQUARE_SIZE // 5
# Getting the eyes' coordinates
if dir in ["U", "D"]:
center1 = (square_top_left[0] + SQUARE_SIZE//4 - 1, square_top_left[1] + SQUARE_SIZE//2)
center2 = (square_top_left[0] + 3 * SQUARE_SIZE//4, square_top_left[1] + SQUARE_SIZE//2)
if dir == "U": # Looking up
center1pupil = (center1[0], center1[1] - radius//2)
center2pupil = (center2[0], center2[1] - radius//2)
elif dir == "D": # Looking down
center1pupil = (center1[0], center1[1] + radius//2)
center2pupil = (center2[0], center2[1] + radius//2)
elif dir in ["R", "L"]:
center1 = (square_top_left[0] + SQUARE_SIZE//2, square_top_left[1] + SQUARE_SIZE//4 - 1)
center2 = (square_top_left[0] + SQUARE_SIZE//2, square_top_left[1] + 3 * SQUARE_SIZE//4)
if dir == "R": # Looking right
center1pupil = (center1[0] + radius//2, center1[1])
center2pupil = (center2[0] + radius//2, center2[1])
elif dir == "L": # Looking left
center1pupil = (center1[0] - radius//2, center1[1])
center2pupil = (center2[0] - radius//2, center2[1])
# Drawing the eyes
cv2.circle(image, center1, radius, (200, 200, 200), -1)
cv2.circle(image, center2, radius, (200, 200, 200), -1)
cv2.circle(image, center1pupil, radius//2, (20, 20, 20), -1)
cv2.circle(image, center2pupil, radius//2, (20, 20, 20), -1)
# Drawing Snake's 1 (green) parts
for pos in snake1_body[1:]:
draw_pos(pos, SNAKE1_COLOR, False, snake1_dir)
draw_pos(snake1_body[0], SNAKE1_COLOR, True, snake1_dir)
# Drawing Snake's 2 (blue) parts
for pos in snake2_body[1:]:
draw_pos(pos, SNAKE2_COLOR, False, snake2_dir)
draw_pos(snake2_body[0], SNAKE2_COLOR, True, snake2_dir)
# Drawing food
for pos in food:
draw_pos(pos, FOOD_COLOR)
image = image.astype(np.uint8)
# Display the image on your computer screen (for testing)
if is_display:
cv2.imshow("Grid Image", image)
cv2.waitKey(0)
cv2.destroyAllWindows()
# Save the image as a PNG file (not needed for the app)
if save_dir is not None:
# Save the image as a PNG file
cv2.imwrite(save_dir + f"{turn:03}_turn.png", image)
else:
return image
if __name__ == "__main__":
# --- Test case ---
# Configs
board_config = {
"GRID_SIZE": 15,
"SQUARE_SIZE": 35,
"LINE_THICKNESS": 2,
"BACKGROUND_COLOR": (30, 20, 20),
"LINES_COLOR": (75, 40, 40),
"SNAKE1_COLOR": (20, 200, 20),
"SNAKE2_COLOR": (190, 120, 0),
"FOOD_COLOR": (50, 50, 250),
"MAX_TURNS": 100,
}
# Initial state
board_state = {
"turn": 0,
"snake1": {
"body": [(5, 2), (4, 2), (3, 2), (2, 2)],
"dir": "R",
"is_alive": True,
},
"snake2": {
"body": [(9, 12), (10, 12), (11, 12), (12, 12)],
"dir": "L",
"is_alive": True,
},
"food": [],
}
board_plot(board_config, board_state, is_display=True, save_dir="./")