89 lines
2.4 KiB
Python
89 lines
2.4 KiB
Python
import time
|
|
|
|
from models.detection.detector import *
|
|
from models.detection.board_manager import *
|
|
from models.detection.pieces_manager import *
|
|
|
|
def print_board_console(board_array):
|
|
files = "abcdefgh"
|
|
print(" " + " ".join(files))
|
|
print(" +-----------------+")
|
|
|
|
for r in range(8):
|
|
row_str = str(8 - r) + "|"
|
|
for f in range(8):
|
|
piece = board_array[r][f]
|
|
if piece is None:
|
|
row_str += ". "
|
|
else:
|
|
if piece[0] == "w":
|
|
piece = piece.upper()
|
|
row_str += piece[2] + " "
|
|
row_str += "|"
|
|
print(row_str)
|
|
print(" +-----------------+")
|
|
|
|
if __name__ == "__main__":
|
|
|
|
print("Initializing models...")
|
|
edges_detector = Detector("../assets/models/edges.pt")
|
|
pieces_detector = Detector("../assets/models/unified-nano-refined.pt")
|
|
|
|
pieces_manager = PiecesManager()
|
|
board_manager = BoardManager()
|
|
|
|
print("Initializing camera...")
|
|
cap = cv2.VideoCapture(0)
|
|
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
|
|
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
|
|
cap.set(cv2.CAP_PROP_FPS, 30)
|
|
|
|
print("Initialized")
|
|
if not cap.isOpened():
|
|
print("Error: Could not open camera")
|
|
exit()
|
|
|
|
cv2.namedWindow("Predictions", cv2.WINDOW_NORMAL)
|
|
|
|
while True:
|
|
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print("Error: Failed to grab frame")
|
|
break
|
|
|
|
edges_pred = edges_detector.make_prediction(frame)
|
|
pieces_pred = pieces_detector.make_prediction(frame)
|
|
|
|
edges_annotated_frame = edges_pred[0].plot()
|
|
pieces_annotated_frame = pieces_pred[0].plot(img=edges_annotated_frame)
|
|
|
|
cv2.imshow("Predictions", pieces_annotated_frame)
|
|
cv2.resizeWindow("Predictions", 640, 640)
|
|
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
remap_width = 800
|
|
remap_height = 800
|
|
|
|
board_manager = BoardManager()
|
|
result = board_manager.process_frame(edges_pred[0], frame, (remap_width, remap_height))
|
|
if result is None:
|
|
continue
|
|
|
|
corners = result[0]
|
|
matrix = result[1]
|
|
|
|
detections = pieces_manager.extract_pieces(pieces_pred)
|
|
|
|
board = pieces_manager.pieces_to_board(detections, matrix, (remap_width, remap_height))
|
|
|
|
fen = pieces_manager.board_to_fen(board)
|
|
print("FEN:", fen)
|
|
|
|
print_board_console(board)
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|