snake-solver/main.py
2025-04-04 08:39:02 +01:00

397 lines
11 KiB
Python

import os
from subprocess import Popen, PIPE
import numpy as np
import cv2
import threading as th
import time as t
import hyprpy as hypr
from evdev import UInput, ecodes as e
import math
# Function to create and initialize a virtual keyboard
def create_virtual_keyboard():
capabilities = {
e.EV_KEY: [e.KEY_W, e.KEY_A, e.KEY_D, e.KEY_S]
}
return UInput(capabilities)
virtual_keyboard = create_virtual_keyboard()
print(virtual_keyboard)
# Function to send a keystroke
def send_keystroke(virtual_keyboard, key):
virtual_keyboard.write(e.EV_KEY, key, 1) # Key down
t.sleep(0.01)
virtual_keyboard.write(e.EV_KEY, key, 0) # Key up
virtual_keyboard.syn()
instance = hypr.Hyprland()
print("please select the game area")
stream = os.popen('slurp')
capture_area = stream.read()
size = list(map(int, capture_area.split(' ')[1].split('x')))
boxPos = list(map(int, capture_area.split(' ')[0].split(',')))
print(capture_area, size)
height = 0
width = 0
def colorCodeToColorArary(target: str):
return np.array([int(target[4:6], 16), int(target[2:4], 16), int(target[0:2], 16)])
green_1 = colorCodeToColorArary("AAD751")
green_2 = colorCodeToColorArary("A2D149")
fruit_color_1 = colorCodeToColorArary("D8B81C")
fruit_color_2 = colorCodeToColorArary("CCAD1E")
def get_img():
global height, width
stream = Popen(['grim', '-g', capture_area.strip(), '-'], stdout=PIPE)
img = cv2.imdecode(np.frombuffer(stream.stdout.read(), dtype='uint8'), cv2.IMREAD_COLOR)
height, width, _ = img.shape
return img
def get_start(img):
start = None
for i in range(height):
for j in range(width):
if start is None:
if (img[i, j] == green_1).all():
print("found target at start", i, j)
start = [i, j]
else:
if (img[i, j] == green_2).all():
print("found end at", i, j)
return [*start, j - start[1]]
if start is not None:
print("Something is wrong found start but not end at the same line")
exit(1)
img = get_img()
cv2.imshow('image', img)
cv2.waitKey(1)
start = get_start(img)
targets = []
# 0 - empty
# 1 - fruit
# 2 - snake
# 3 - head
targets_v = {}
head = None
food = None
snake_part = []
for i in range(15):
for j in range(17):
i2 = i * start[2] + int(start[2] / 2) + start[0]
j2 = j * start[2] + int(start[2] / 2) + start[1]
targets.append([i2, j2, i, j])
targets_v[f"{i}-{j}"] = None
def draw_square(img, pos, color):
i, j = pos
i2 = int(i * start[2] + int(start[2] / 2) + start[0])
j2 = int(j * start[2] + int(start[2] / 2) + start[1])
img[i2-5:i2+5, j2-5:j2+5] = color
def draw_dots_from_data(img):
global food, head
for target in targets:
i2, j2, i, j = target
v = targets_v[f"{i}-{j}"]
if v == 0:
# img[i2-5:i2+5, j2-5:j2+5] = [0, 255, 0]
pass
elif v == 1:
img[i2-5:i2+5, j2-5:j2+5] = [0, 0, 255]
elif v == 2:
img[i2-5:i2+5, j2-5:j2+5] = [255, 0, 0]
elif v == 3:
img[i2-5:i2+5, j2-5:j2+5] = [255, 255, 255]
else:
img[i2-5:i2+5, j2-5:j2+5] = [255, 0, 255]
def draw_dots(img):
global food, head, snake_part, head_p
snake_part = []
for target in targets:
i2, j2, i, j = target
color = img[i2, j2]
# body samples
bs_1 = img[i2 - 7, j2]
bs_2 = img[i2 + 7, j2]
bs_3 = img[i2, j2 - 7]
bs_4 = img[i2, j2 + 7]
all_the_same = (bs_1 == bs_2).all() and (bs_1 == bs_3).all() and (bs_1 == bs_4).all()
bs_list = np.array([bs_1, bs_2, bs_3, bs_4])
if (color == green_1).all() or (color == green_2).all() and all_the_same:
# img[i2-5:i2+5, j2-5:j2+5] = [0, 255, 0]
targets_v[f"{i}-{j}"] = 0
elif (color == fruit_color_1).all() or (color == fruit_color_2).all():
img[i2-5:i2+5, j2-5:j2+5] = [0, 0, 255]
targets_v[f"{i}-{j}"] = 1
food = [i, j]
elif color[0] > 150:
img[i2-5:i2+5, j2-5:j2+5] = [255, 0, 0]
targets_v[f"{i}-{j}"] = 2
snake_part.append([i, j])
else:
greens_1 = np.sum(bs_list == green_1, axis=1)
greens_2 = np.sum(bs_list == green_2, axis=1)
if np.sum(greens_1) == 9 and bs_list[np.argmin(greens_1)][0] > 150:
targets_v[f"{i}-{j}"] = 3
img[i2-5:i2+5, j2-5:j2+5] = [255, 255, 255]
head = np.array([i, j])
if np.sum(greens_2) == 9 and bs_list[np.argmin(greens_2)][0] > 150:
targets_v[f"{i}-{j}"] = 3
img[i2-5:i2+5, j2-5:j2+5] = [255, 255, 255]
head = np.array([i, j])
else:
img[i2, j2] = [255, 0, 255]
img[i2 - 7, j2] = [255, 0, 255]
img[i2 + 7, j2] = [255, 0, 255]
img[i2, j2 - 7] = [255, 0, 255]
img[i2, j2 + 7] = [255, 0, 255]
targets_v[f"{i}-{j}"] = None
def pt_dist(a, b):
return math.sqrt((a[0] - b[0])**2 + (a[1] - b[1])**2)
def get_all_adj(a):
return [
[a[0] + 1, a[1]],
[a[0] - 1, a[1]],
[a[0], a[1] + 1],
[a[0], a[1] - 1],
]
def draw_path(img, ps, final=False):
for p in ps:
i, j = p
i2 = i * start[2] + int(start[2] / 2) + start[0]
j2 = j * start[2] + int(start[2] / 2) + start[1]
if final:
img[i2-3:i2+3, j2-3:j2+3] = [255, 255, 0]
else:
img[i2-3:i2+3, j2-3:j2+3] = [0, 255, 255]
def calculate_path(pos, visited, img=None):
# if img is not None and visited is not None:
# draw_dots_from_data(img)
# draw_path(img, visited)
# cv2.imshow('image-thread', img)
# cv2.waitKey(1)
n_visited = visited.copy()
n_visited.append(pos)
queue = get_all_adj(pos)
queue.sort(key=lambda p: pt_dist(food, p), reverse=False)
for i in queue:
if i in n_visited:
# print("already visisted", i)
continue
v = targets_v.get(f'{i[0]}-{i[1]}')
if v == 1:
# print("is food", i)
return [pos, i]
if v != 0:
# print("is ocupied", i)
continue
maybe_path = calculate_path(i, n_visited, img=img)
if maybe_path is not None:
maybe_path.insert(0, pos)
return maybe_path
return None
draw_dots(img)
cv2.imshow('image', img)
cv2.waitKey(1)
# Good for testing
# cv2.waitKey(0)
# exit(1)
last = t.time_ns()
g_img = img
def img_collection_and_path_analisis(data):
print("start img col")
path = None
while True:
img = get_img()
data['img'] = img
draw_dots(img)
if path is not None:
draw_path(img, path, final=True)
if data['path'] is None:
if head is None:
print("could not find head skiping")
t.sleep(0.001)
continue
if food is None:
print("could not find food skiping")
t.sleep(0.001)
continue
# print("calculate path")
maybe_path = calculate_path(list(head), [], img)
if maybe_path is None:
# print("No path found")
pass
draw_dots_from_data(img)
if maybe_path is not None:
draw_path(img, maybe_path, True)
# print("got_path")
path = maybe_path
data['path'] = maybe_path
data = {'path': None, 'img': None}
img_col = th.Thread(target=img_collection_and_path_analisis, args=(data, ))
img_col.start()
last = t.time_ns()
head_pred = None
time_to_pred = None
start_time_to_pred = None
last_time_to_pred = 0
lp_avg = 0
lp_count = 0
while True:
# sleep mode if the mouse is outside the area
pos = instance.get_cursor_pos()
if not ((pos[0] > boxPos[0]) and (pos[0] < boxPos[0] + size[0]) and (pos[1] > boxPos[1]) and (pos[1] < boxPos[1] + size[1])):
t.sleep(0.5)
cv2.waitKey(1)
continue
now = t.time_ns()
lp_count += 1
lp_time = (now - last)/1000000
lp_avg += lp_time
time_since_last_move = (now - start_time_to_pred)/1000000 if start_time_to_pred is not None else 0
lp_time_next = (last_time_to_pred - time_since_last_move) / (lp_avg/lp_count)
print(
f'\r lp time {lp_time:.4f}ms lp avg {lp_avg/lp_count:.4f}ms',
"pred: ", last_time_to_pred,
f"loops til next: {lp_time_next:.2f}",
f"loops time since last move {time_since_last_move:.2f}",
" ms ", end="")
last = t.time_ns()
path = data['path']
img = data['img'].copy()
if img is not None:
cv2.imshow('image', img)
cv2.waitKey(1)
l_head = list(head)
if l_head == head_pred and time_to_pred is None:
time_to_pred = t.time_ns() - start_time_to_pred
last_time_to_pred = (t.time_ns() - start_time_to_pred)/1000000
print("pred hit in ", time_to_pred / 1000000, "ms!", end=" ")
if head is not None and path is not None:
if l_head in path:
index = path.index(l_head)
if index < len(path) - 1:
diff = np.array(l_head) - np.array(path[index + 1])
if diff[1] == 1:
# print('turn Left', end="")
send_keystroke(virtual_keyboard, e.KEY_A)
elif diff[1] == -1:
# print('turn Right', end="")
send_keystroke(virtual_keyboard, e.KEY_D)
elif diff[0] == 1:
# print('turn Up', end="")
send_keystroke(virtual_keyboard, e.KEY_W)
elif diff[0] == -1:
# print('turn Down', end="")
send_keystroke(virtual_keyboard, e.KEY_S)
else:
print("not sure what to do", diff, end="")
exit(1)
if head_pred != path[index + 1]:
head_pred = path[index + 1]
start_time_to_pred = t.time_ns()
time_to_pred = None
else:
print('head not in path', end="")
data['path'] = None
else:
# print("head", head, "path", path)
pass
print(" ", end="")
cv2.waitKey(0)
# import pygame
# scrn = pygame.display.set_mode(size)
# running = True
# while (running):
# for i in pygame.event.get():
# if i.type == pygame.QUIT:
# running = False
#
# stream = Popen(['grim', '-g', capture_area.strip(), '-'], stdout=PIPE)
# # print(stream.stdout.read())
#
# # img = pygame.image.load(stream.stdout, "test.png")
#
# print(len(stream.stdout.read()))
#
# # img = pygame.image.frombuffer(stream.stdout.read(), size, "RGB")
#
# # scrn.blit(img, (0, 0))
#
# # pygame.display.flip()
#
# pygame.quit()