143 lines
6.5 KiB (Stored with Git LFS)
Python
143 lines
6.5 KiB (Stored with Git LFS)
Python
import cv2
|
|
import argparse
|
|
import enum
|
|
import abc
|
|
import numpy as np
|
|
|
|
from board_detction import Board_19x19, VISION_PREPROCESSING_MODE, Board_Detection_SM, detect_board
|
|
|
|
# Parse arguments
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("camera", type=int, help="Camera index")
|
|
args = parser.parse_args()
|
|
|
|
cam = cv2.VideoCapture(args.camera)
|
|
|
|
WINDOW_NAME = "Output Go Board Detection"
|
|
cv2.namedWindow(WINDOW_NAME)
|
|
cv2.moveWindow(WINDOW_NAME, 1,0)
|
|
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
font_color = (127, 255, 0)
|
|
|
|
fps = 0
|
|
|
|
board_messurements = Board_19x19()
|
|
state_machine = Board_Detection_SM()
|
|
|
|
try:
|
|
while True:
|
|
t_start = cv2.getTickCount()
|
|
ret, frame = cam.read()
|
|
width, height = frame.shape[1], frame.shape[0]
|
|
|
|
frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
frame_edge = cv2.GaussianBlur(frame_gray, (5, 5), 0)
|
|
frame_edge = cv2.adaptiveThreshold(frame_edge, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 9, 2)
|
|
frame_edge_copy = cv2.cvtColor(frame_edge.copy(), cv2.COLOR_GRAY2BGR)
|
|
|
|
board_contures_candiates, frame_proc = detect_board(state_machine.state, frame)
|
|
|
|
board_conture = state_machine.advance_state(board_contures_candiates)
|
|
|
|
frame_dection = cv2.drawContours(frame.copy(), board_contures_candiates, -1, (0, 255, 0), 1)
|
|
|
|
frame_classifications = frame.copy()
|
|
|
|
if board_conture is not None:
|
|
N_line_positions = []
|
|
S_line_positions = []
|
|
W_line_positions = []
|
|
E_line_positions = []
|
|
|
|
for i in range(19):
|
|
N_line_positions.append(board_messurements.get_cell_position_on_line(i, board_conture[0], board_conture[1]))
|
|
E_line_positions.append(board_messurements.get_cell_position_on_line(i, board_conture[1], board_conture[2]))
|
|
S_line_positions.append(board_messurements.get_cell_position_on_line(i, board_conture[2], board_conture[3]))
|
|
W_line_positions.append(board_messurements.get_cell_position_on_line(i, board_conture[3], board_conture[0]))
|
|
|
|
N_line_positions = np.array(N_line_positions, dtype=np.int32)
|
|
E_line_positions = np.array(E_line_positions, dtype=np.int32)
|
|
S_line_positions = np.array(list(reversed(S_line_positions)), dtype=np.int32)
|
|
W_line_positions = np.array(list(reversed(W_line_positions)), dtype=np.int32)
|
|
|
|
cv2.circle(frame_dection, board_conture[0].astype(np.int32), 3, (0, 255, 255), -1)
|
|
cv2.circle(frame_dection, board_conture[1].astype(np.int32), 3, (0, 255, 255), -1)
|
|
cv2.circle(frame_dection, board_conture[2].astype(np.int32), 3, (0, 255, 255), -1)
|
|
cv2.circle(frame_dection, board_conture[3].astype(np.int32), 3, (0, 255, 255), -1)
|
|
|
|
for p in N_line_positions:
|
|
cv2.circle(frame_dection, p, 2, (0, 0, 255), -1)
|
|
for p in E_line_positions:
|
|
cv2.circle(frame_dection, p, 2, (0, 0, 255), -1)
|
|
for p in S_line_positions:
|
|
cv2.circle(frame_dection, p, 2, (0, 0, 255), -1)
|
|
for p in W_line_positions:
|
|
cv2.circle(frame_dection, p, 2, (0, 0, 255), -1)
|
|
|
|
for p0, p1 in zip(N_line_positions, S_line_positions):
|
|
cv2.line(frame_dection, tuple(p0), tuple(p1), (0, 0, 255), 1)
|
|
|
|
for i in range(19):
|
|
p2 = board_messurements.get_cell_position_on_line(i, p0, p1)
|
|
cv2.circle(frame_dection, p2, 2, (0, 0, 255), -1)
|
|
|
|
for p0, p1 in zip(E_line_positions, W_line_positions):
|
|
cv2.line(frame_dection, tuple(p0), tuple(p1), (0, 255, 0), 1)
|
|
|
|
for i in range(19):
|
|
p2 = board_messurements.get_cell_position_on_line(i, p0, p1)
|
|
cv2.circle(frame_dection, p2, 2, (0, 255, 0), -1)
|
|
|
|
MESSURE_WIDTH = 15
|
|
STONE_DECTION_WIDTH = 5
|
|
for x in range(19):
|
|
for y in range(19):
|
|
mid_point = board_messurements.get_cell_position_on_line(x, N_line_positions[y], S_line_positions[y])
|
|
subsection = frame[mid_point[1] - MESSURE_WIDTH//2:mid_point[1] + MESSURE_WIDTH//2, mid_point[0] - MESSURE_WIDTH//2:mid_point[0] + MESSURE_WIDTH//2]
|
|
avrage_color = np.mean(subsection, axis=(0, 1)).astype(np.int32).tolist()
|
|
|
|
cv2.rectangle(frame_classifications, (mid_point[0] - MESSURE_WIDTH//2, mid_point[1] - MESSURE_WIDTH//2), (mid_point[0] + MESSURE_WIDTH//2, mid_point[1] + MESSURE_WIDTH//2), avrage_color, -1)
|
|
cv2.rectangle(frame_classifications, (mid_point[0] - MESSURE_WIDTH//2, mid_point[1] - MESSURE_WIDTH//2), (mid_point[0] + MESSURE_WIDTH//2, mid_point[1] + MESSURE_WIDTH//2), (0, 0, 255), 1)
|
|
|
|
avrage_stone_color = np.mean(frame_edge, axis=(0, 1)).astype(np.int32).tolist()
|
|
cv2.rectangle(frame_edge_copy, (mid_point[0] - STONE_DECTION_WIDTH//2, mid_point[1] - STONE_DECTION_WIDTH//2), (mid_point[0] + STONE_DECTION_WIDTH//2, mid_point[1] + STONE_DECTION_WIDTH//2), (avrage_stone_color, avrage_stone_color, avrage_stone_color), -1)
|
|
cv2.rectangle(frame_edge_copy, (mid_point[0] - STONE_DECTION_WIDTH//2, mid_point[1] - STONE_DECTION_WIDTH//2), (mid_point[0] + STONE_DECTION_WIDTH//2, mid_point[1] + STONE_DECTION_WIDTH//2), (0, 0, 255), 1)
|
|
|
|
res = np.concatenate((
|
|
np.concatenate((
|
|
frame,
|
|
frame_classifications,
|
|
frame_edge_copy
|
|
), axis=1),
|
|
np.concatenate((
|
|
cv2.cvtColor(frame_proc, cv2.COLOR_GRAY2BGR),
|
|
frame_dection,
|
|
frame
|
|
), axis=1)
|
|
), axis=0)
|
|
|
|
for n,t in enumerate([
|
|
f"FPS: {fps:.2f}",
|
|
f"GoBoard Dections: {len(board_contures_candiates)}",
|
|
] + state_machine.create_info_text()):
|
|
res = cv2.putText(res, t, (10, 30 + 15*n), font, 0.5, font_color, 1, cv2.LINE_AA)
|
|
|
|
cv2.imshow(WINDOW_NAME, res)
|
|
|
|
t_end = cv2.getTickCount()
|
|
|
|
fps = cv2.getTickFrequency() / (t_end - t_start)
|
|
|
|
k = cv2.waitKey(1)
|
|
if k%256 == 27: # ESC pressed
|
|
print("Escape hit, closing...")
|
|
break
|
|
|
|
|
|
except KeyboardInterrupt:
|
|
print("Exiting")
|
|
|
|
cam.release()
|
|
cv2.destroyAllWindows()
|