diff options
Diffstat (limited to 'src/argaze/ArUcoMarkers/ArUcoDetector.py')
-rw-r--r-- | src/argaze/ArUcoMarkers/ArUcoDetector.py | 303 |
1 files changed, 303 insertions, 0 deletions
diff --git a/src/argaze/ArUcoMarkers/ArUcoDetector.py b/src/argaze/ArUcoMarkers/ArUcoDetector.py new file mode 100644 index 0000000..86bcbbf --- /dev/null +++ b/src/argaze/ArUcoMarkers/ArUcoDetector.py @@ -0,0 +1,303 @@ +#!/usr/bin/env python + +from typing import TypeVar, Tuple +from dataclasses import dataclass, field +import json +from collections import Counter + +from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoMarker, ArUcoCamera + +import numpy +import cv2 as cv +import cv2.aruco as aruco + +ArUcoMarkerDictionaryType = TypeVar('ArUcoMarkerDictionary', bound="ArUcoMarkerDictionary") +# Type definition for type annotation convenience + +ArUcoMarkerType = TypeVar('ArUcoMarker', bound="ArUcoMarker") +# Type definition for type annotation convenience + +ArUcoCameraType = TypeVar('ArUcoCamera', bound="ArUcoCamera") +# Type definition for type annotation convenience + +DetectorParametersType = TypeVar('DetectorParameters', bound="DetectorParameters") +# Type definition for type annotation convenience + +ArUcoDetectorType = TypeVar('ArUcoDetector', bound="ArUcoDetector") +# Type definition for type annotation convenience + +class DetectorParameters(): + """Define ArUco marker detector parameters. + + .. note:: More details on [opencv page](https://docs.opencv.org/4.x/d1/dcd/structcv_1_1aruco_1_1DetectorParameters.html) + """ + + __parameters = aruco.DetectorParameters_create() + __parameters_names = [ + 'adaptiveThreshConstant', + 'adaptiveThreshWinSizeMax', + 'adaptiveThreshWinSizeMin', + 'adaptiveThreshWinSizeStep', + 'aprilTagCriticalRad', + 'aprilTagDeglitch', + 'aprilTagMaxLineFitMse', + 'aprilTagMaxNmaxima', + 'aprilTagMinClusterPixels', + 'aprilTagMinWhiteBlackDiff', + 'aprilTagQuadDecimate', + 'aprilTagQuadSigma', + 'cornerRefinementMaxIterations', + 'cornerRefinementMethod', + 'cornerRefinementMinAccuracy', + 'cornerRefinementWinSize', + 'markerBorderBits', + 'minMarkerPerimeterRate', + 'maxMarkerPerimeterRate', + 'minMarkerDistanceRate', + 'detectInvertedMarker', + 'errorCorrectionRate', + 'maxErroneousBitsInBorderRate', + 'minCornerDistanceRate', + 'minDistanceToBorder', + 'minOtsuStdDev', + 'perspectiveRemoveIgnoredMarginPerCell', + 'perspectiveRemovePixelPerCell', + 'polygonalApproxAccuracyRate' + ] + + def __init__(self, **kwargs): + + for parameter, value in kwargs.items(): + + setattr(self.__parameters, parameter, value) + + self.__dict__.update(kwargs) + + def __setattr__(self, parameter, value): + + setattr(self.__parameters, parameter, value) + + def __getattr__(self, parameter): + + return getattr(self.__parameters, parameter) + + @classmethod + def from_json(self, json_filepath) -> DetectorParametersType: + """Load detector parameters from .json file.""" + + with open(json_filepath) as configuration_file: + + return DetectorParameters(**json.load(configuration_file)) + + def __str__(self, print_all=False) -> str: + """Detector paremeters string representation.""" + + output = '' + + for parameter in self.__parameters_names: + + if parameter in self.__dict__.keys(): + + output += f'\t*{parameter}: {getattr(self.__parameters, parameter)}\n' + + elif print_all: + + output += f'\t{parameter}: {getattr(self.__parameters, parameter)}\n' + + return output + + @property + def internal(self): + return self.__parameters + +class ArUcoDetector(): + """ArUco markers detector.""" + + dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary = field(init=False, default_factory=ArUcoMarkersDictionary.ArUcoMarkersDictionary) + """ArUco markers dictionary to detect.""" + + marker_size: float = field(init=False) + """Size of ArUco markers to detect in centimeter.""" + + camera: ArUcoCamera.ArUcoCamera = field(init=False, default_factory=ArUcoCamera.ArUcoCamera) + """ArUco camera ...""" + + def __init__(self, **kwargs): + + self.dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(kwargs.pop('dictionary')) + self.marker_size = kwargs.pop('marker_size') + self.camera = ArUcoCamera.ArUcoCamera(**kwargs.pop('camera')) + + # Init detector parameters + self.__parameters = DetectorParameters(**kwargs.pop('parameters')) + + # Init detected markers data + self.__detected_markers = {} + self.__detected_markers_corners = [] + self.__detected_markers_ids = [] + + # Init detected board data + self.__board = None + self.__board_corners_number = 0 + self.__board_corners = [] + self.__board_corners_ids = [] + + # Init detect metrics data + self.__detection_count = 0 + self.__detected_ids = [] + + @classmethod + def from_json(self, json_filepath: str) -> ArUcoDetectorType: + """Load ArUcoDetector setup from .json file.""" + + with open(json_filepath) as configuration_file: + + return ArUcoDetector(**json.load(configuration_file)) + + def __str__(self) -> str: + """String display""" + + output = f'Camera:\n{self.camera}\n' + output += f'Parameters:\n{self.__parameters}\n' + + return output + + @property + def parameters(self): + """ArUco marker detector parameters.""" + + return self.__parameters + + def detect(self, frame): + """Detect all ArUco markers into a frame. + + .. danger:: DON'T MIRROR FRAME + It makes the markers detection to fail. + """ + + # Reset detected markers data + self.__detected_markers, self.__detected_markers_corners, self.__detected_markers_ids = {}, [], [] + + # Detect markers into gray picture + self.__detected_markers_corners, self.__detected_markers_ids, _ = aruco.detectMarkers(cv.cvtColor(frame, cv.COLOR_BGR2GRAY), self.dictionary.markers, parameters = self.__parameters.internal) + + # Is there detected markers ? + if len(self.__detected_markers_corners) > 0: + + # Gather detected markers data and update metrics + self.__detection_count += 1 + + for i, marker_id in enumerate(self.__detected_markers_ids.T[0]): + + marker = ArUcoMarker.ArUcoMarker(self.dictionary, marker_id, self.marker_size) + + marker.corners = self.__detected_markers_corners[i] + + # No pose estimation: call estimate_pose to get one + marker.translation = numpy.empty([0]) + marker.rotation = numpy.empty([0]) + marker.points = numpy.empty([0]) + + self.__detected_markers[marker_id] = marker + + self.__detected_ids.append(marker_id) + + def estimate_pose(self): + """Estimate pose of current detected markers.""" + + # Is there detected markers ? + if len(self.__detected_markers_corners) > 0: + + markers_rvecs, markers_tvecs, markers_points = aruco.estimatePoseSingleMarkers(self.__detected_markers_corners, self.marker_size, numpy.array(self.camera.K), numpy.array(self.camera.D)) + + for i, marker_id in enumerate(self.__detected_markers_ids.T[0]): + + marker = self.__detected_markers[marker_id] + + marker.translation = markers_tvecs[i][0] + marker.rotation, _ = cv.Rodrigues(markers_rvecs[i][0]) + marker.points = markers_points.reshape(4, 3) + + @property + def detected_markers(self) -> dict[ArUcoMarkerType]: + """Access to detected markers dictionary.""" + + return self.__detected_markers + + @property + def detected_markers_number(self) -> int: + """Return detected markers number.""" + + return len(list(self.__detected_markers.keys())) + + def draw_detected_markers(self, frame): + """Draw traked markers.""" + + for marker_id, marker in self.__detected_markers.items(): + + marker.draw(frame, self.camera.K, self.camera.D) + + def detect_board(self, frame, board, expected_markers_number): + """Detect ArUco markers board in frame setting up the number of detected markers needed to agree detection. + + .. danger:: DON'T MIRROR FRAME + It makes the markers detection to fail. + """ + + # detect markers from gray picture + gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) + self.__detected_markers_corners, self.__detected_markers_ids, _ = aruco.detectMarkers(gray, self.dictionary.markers, parameters = self.__parameters.internal) + + # if all board markers are detected + if len(self.__detected_markers_corners) == expected_markers_number: + + self.__board = board + self.__board_corners_number, self.__board_corners, self.__board_corners_ids = aruco.interpolateCornersCharuco(self.__detected_markers_corners, self.__detected_markers_ids, gray, self.__board.model) + + else: + + self.__board = None + self.__board_corners_number = 0 + self.__board_corners = [] + self.__board_corners_ids = [] + + def draw_board(self, frame): + """Draw detected board corners in frame.""" + + if self.__board != None: + + cv.drawChessboardCorners(frame, ((self.__board.size[0] - 1 ), (self.__board.size[1] - 1)), self.__board_corners, True) + + def reset_detection_metrics(self): + """Enable marker detection metrics.""" + + self.__detection_count = 0 + self.__detected_ids = [] + + @property + def detection_metrics(self) -> Tuple[int, dict]: + """Get marker detection metrics. + * **Returns:** + - number of detect function call + - dict with number of detection for each marker identifier""" + + return self.__detection_count, Counter(self.__detected_ids) + + @property + def board_corners_number(self) -> int: + """Get detected board corners number.""" + + return self.__board_corners_number + + @property + def board_corners_identifier(self) -> list[int]: + """Get detected board corners identifier.""" + + return self.__board_corners_ids + + @property + def board_corners(self) -> list: + """Get detected board corners.""" + + return self.__board_corners + |