101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
import cv2
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
from ultralytics.engine.results import Results
|
|
|
|
from hardware.camera.camera import Camera
|
|
from models.detection.detector import Detector
|
|
from models.detection.board_manager import BoardManager
|
|
from models.detection.pieces_manager import PiecesManager
|
|
|
|
|
|
class DetectionService:
|
|
|
|
edges_detector : Detector
|
|
pieces_detector : Detector
|
|
|
|
board_manager : BoardManager
|
|
pieces_manager : PiecesManager
|
|
|
|
scale_size : tuple[int, int]
|
|
|
|
camera : Camera
|
|
|
|
def __init__(self):
|
|
current_file = Path(__file__).resolve()
|
|
project_root = current_file.parent.parent
|
|
|
|
self.edges_detector = Detector(project_root / "assets" / "models" / "edges.pt")
|
|
self.pieces_detector = Detector(project_root / "assets" / "models" / "unified-nano-refined.pt")
|
|
|
|
self.pieces_manager = PiecesManager()
|
|
self.board_manager = BoardManager()
|
|
self.scale_size = (800, 800)
|
|
self.camera = Camera()
|
|
|
|
def start(self):
|
|
self.camera.open()
|
|
|
|
def stop(self):
|
|
self.camera.close()
|
|
|
|
def analyze_single_frame(self) -> tuple[np.ndarray, str | None]:
|
|
frame = self.camera.take_photo()
|
|
fen = self.__get_fen(frame)
|
|
return frame, fen
|
|
|
|
|
|
def __run_complete_detection(self, frame : np.ndarray, display=False) -> dict[str, list[Results]] :
|
|
pieces_prediction = self.__run_pieces_detection(frame)
|
|
edges_prediction = self.__run_edges_detection(frame)
|
|
|
|
if display:
|
|
edges_annotated_frame = edges_prediction[0].plot()
|
|
pieces_annotated_frame = pieces_prediction[0].plot(img=edges_annotated_frame)
|
|
self.__display_frame(pieces_annotated_frame)
|
|
|
|
return { "edges" : edges_prediction, "pieces" : pieces_prediction}
|
|
|
|
|
|
def __run_pieces_detection(self, frame : np.ndarray, display=False) -> list[Results]:
|
|
prediction = self.pieces_detector.make_prediction(frame)
|
|
if display:
|
|
self.__display_frame(prediction[0].plot())
|
|
return prediction
|
|
|
|
|
|
def __run_edges_detection(self, frame : np.ndarray, display=False) -> list[Results]:
|
|
prediction = self.edges_detector.make_prediction(frame)
|
|
if display:
|
|
self.__display_frame(prediction[0].plot())
|
|
return prediction
|
|
|
|
|
|
def __get_fen(self, frame : np.ndarray) -> str | None:
|
|
result = self.__run_complete_detection(frame)
|
|
|
|
edges_prediction = result["edges"]
|
|
pieces_prediction = result["pieces"]
|
|
|
|
processed_frame = self.board_manager.process_frame(edges_prediction[0], frame, self.scale_size)
|
|
if processed_frame is None:
|
|
return None
|
|
|
|
warped_corners = processed_frame[0]
|
|
matrix = processed_frame[1]
|
|
|
|
detections = self.pieces_manager.extract_pieces(pieces_prediction)
|
|
|
|
board = self.pieces_manager.pieces_to_board(detections, warped_corners, matrix, self.scale_size)
|
|
|
|
return self.pieces_manager.board_to_fen(board)
|
|
|
|
def __display_frame(self, frame : np.ndarray):
|
|
cv2.namedWindow("Frame", cv2.WINDOW_NORMAL)
|
|
cv2.resizeWindow("Frame", self.scale_size[0], self.scale_size[1])
|
|
cv2.imshow("Frame", frame)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|
|
return
|