160 lines
6.3 KiB (Stored with Git LFS)
Python
160 lines
6.3 KiB (Stored with Git LFS)
Python
import cv2
|
|
import argparse
|
|
import enum
|
|
import uuid
|
|
import numpy as np
|
|
|
|
from model import load_model, TARGET_SIZE, classify_stone_empty
|
|
|
|
|
|
# Parse arguments
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("camera", type=int, help="Camera index")
|
|
args = parser.parse_args()
|
|
|
|
cam = cv2.VideoCapture(args.camera)
|
|
|
|
stone_detection_model = load_model("models/model1.pth")
|
|
|
|
WINDOW_NAME = "Output Go Board Detection"
|
|
cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
|
|
cv2.resizeWindow(WINDOW_NAME, 1920, 1080)
|
|
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()
|
|
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
|
recording = False
|
|
|
|
try:
|
|
while True:
|
|
t_start = cv2.getTickCount()
|
|
ret, frame = cam.read()
|
|
|
|
|
|
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_edge_processed = detect_board(state_machine.state, frame)
|
|
|
|
#board_conture = state_machine.advance_state(board_contures_candiates)
|
|
|
|
|
|
if len(board_contures_candiates) > 0:
|
|
board_conture = np.squeeze(board_contures_candiates[0], axis=1)
|
|
frame_dection = cv2.drawContours(frame.copy(), board_contures_candiates, -1, (0, 255, 0), 1)
|
|
frame_classifications = frame.copy()
|
|
else:
|
|
board_conture = None
|
|
frame_dection = frame.copy()
|
|
frame_classifications = frame.copy()
|
|
|
|
# Stage 2: Classfiy Stones
|
|
if board_conture is not None:
|
|
|
|
# Getting Perspectiv correct Image
|
|
board_image = corret_board_perspective(frame, board_conture, s=400)
|
|
display_board_image = board_image.copy()
|
|
draw_board_markers(display_board_image, np.array([
|
|
[0, 0],
|
|
[0, board_image.shape[1]],
|
|
[board_image.shape[0], board_image.shape[1]],
|
|
[board_image.shape[0], 0]
|
|
]), board_messurements)
|
|
|
|
display_board_image = pad_center(display_board_image, frame.shape[0:2])
|
|
board_masked, board_antimasked1, board_antimasked2, board_pre_processed = preprocess_for_stone_dection(board_image)
|
|
|
|
classification_output = board_image.copy()
|
|
classification_color_output = board_image.copy()
|
|
|
|
pre_classifcation = classfy_board(board_masked, board_antimasked1, board_antimasked2, classification_output, board_messurements, windows_size=10)
|
|
color_classifcation = classfiy_color(board_image, pre_classifcation, classification_color_output, board_messurements, windows_size=10)
|
|
|
|
display_classifcation_preprocessed = pad_center(cv2.cvtColor(board_pre_processed, cv2.COLOR_GRAY2BGR), frame.shape[0:2])
|
|
display_classifcation_output = pad_center(classification_output, frame.shape[0:2])
|
|
display_classifcation_color_output = pad_center(classification_color_output, frame.shape[0:2])
|
|
|
|
else:
|
|
display_board_image = np.full(frame.shape, 0, dtype=np.uint8)
|
|
display_classifcation_preprocessed = np.full(frame.shape, 0, dtype=np.uint8)
|
|
display_classifcation_output = np.full(frame.shape, 0, dtype=np.uint8)
|
|
display_classifcation_color_output = np.full(frame.shape, 0, dtype=np.uint8)
|
|
display_classifcation_preprocessed = np.full(frame.shape, 0, dtype=np.uint8)
|
|
|
|
if recording:
|
|
VideoOut.write(frame)
|
|
|
|
# Added Labeling
|
|
frame_edge_processed = cv2.cvtColor(frame_edge_processed, cv2.COLOR_GRAY2BGR)
|
|
cv2.putText(frame, "Raw Input Frame", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)
|
|
cv2.putText(frame_edge_processed, "Board Edges", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)
|
|
cv2.putText(display_board_image, "Board (Perspecive Corrected)", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)
|
|
cv2.putText(display_classifcation_output, "Stone Pre-Classfication", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)
|
|
cv2.putText(display_classifcation_color_output, "Stone Classfication", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)
|
|
cv2.putText(display_classifcation_preprocessed, "Classfication Preprocessed", (10, frame.shape[0]-10), font, 1, font_color, 2, cv2.LINE_AA)
|
|
|
|
res = np.concatenate((
|
|
np.concatenate((
|
|
frame,
|
|
display_board_image,
|
|
display_classifcation_preprocessed,
|
|
), axis=1),
|
|
np.concatenate((
|
|
frame_edge_processed,
|
|
display_classifcation_output,
|
|
display_classifcation_color_output,
|
|
), axis=1),
|
|
np.concatenate((
|
|
frame_edge_processed,
|
|
display_classifcation_output,
|
|
display_classifcation_color_output,
|
|
), 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)
|
|
|
|
if recording:
|
|
cv2.circle(res, (20, 20), 10, (0, 0, 255), -1)
|
|
|
|
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...")
|
|
break0
|
|
elif k%256 == 32:
|
|
cv2.imwrite("frame.png", frame)
|
|
|
|
elif k%256 == 114:# # R-Key
|
|
if recording:
|
|
recording = False
|
|
VideoOut.release()
|
|
else:
|
|
recording = True
|
|
VideoOut = cv2.VideoWriter(f"vid_{str(uuid.uuid4())}.mp4", fourcc, 30.0, (frame.shape[1], frame.shape[0]))
|
|
|
|
|
|
except KeyboardInterrupt:
|
|
print("Exiting")
|
|
|
|
cam.release()
|
|
cv2.destroyAllWindows()
|