Skip to content
Snippets Groups Projects
Commit 955fa24b authored by Adrien Lescourt's avatar Adrien Lescourt
Browse files

Few tweaks

parent 3d9c8320
No related branches found
No related tags found
No related merge requests found
...@@ -31,6 +31,10 @@ class PieceStateExtractor(ABC): ...@@ -31,6 +31,10 @@ class PieceStateExtractor(ABC):
def train(self, img: np.ndarray) -> None: def train(self, img: np.ndarray) -> None:
... ...
@abstractmethod
def is_ready(self) -> bool:
...
class LightnessPieceStateDetector(PieceStateDetector): class LightnessPieceStateDetector(PieceStateDetector):
""" Primitive detector based on lightness of average pixels around the piece""" """ Primitive detector based on lightness of average pixels around the piece"""
...@@ -265,7 +269,12 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor): ...@@ -265,7 +269,12 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor):
if marker_id in (self.MARKER_ON_BOARD_ID, self.MARKER_OFF_BOARD_ID): if marker_id in (self.MARKER_ON_BOARD_ID, self.MARKER_OFF_BOARD_ID):
r = ((int(x), int(y)) for x, y in corners[idx][0].tolist()) r = ((int(x), int(y)) for x, y in corners[idx][0].tolist())
rect = Rect.from_corners(*(Point.from_tuple(t) for t in r)) rect = Rect.from_corners(*(Point.from_tuple(t) for t in r))
# TODO
# BUG: the center is only correct if the marker is in the right orientation
# otherwise, we get complute the center on a corner
center = get_rect_center(rect) center = get_rect_center(rect)
cv.circle(img, (center.x, center.y), 10, (123,0,0))
cv.rectangle(img, (rect.top_left.x, rect.top_left.y), (rect.top_left.x+rect.width, rect.top_left.y+rect.height), (255, 0 ,0))
all_markers.append(center) all_markers.append(center)
if len(all_markers) == 16: if len(all_markers) == 16:
...@@ -287,7 +296,7 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor): ...@@ -287,7 +296,7 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor):
) )
# DEBUG # DEBUG
cv.aruco.drawDetectedMarkers(img, corners) # cv.aruco.drawDetectedMarkers(img, corners)
if ids is not None: if ids is not None:
all_markers = [] all_markers = []
...@@ -298,7 +307,7 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor): ...@@ -298,7 +307,7 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor):
all_markers.append((marker_ids[0], center)) all_markers.append((marker_ids[0], center))
for marker_id, center in all_markers: for marker_id, center in all_markers:
closest_piece_idx = np.argmin( closest_piece_idx = np.argmin( # argmin to get the index or the min
[get_segment_size(center, piece.postion) for piece in board.pieces] [get_segment_size(center, piece.postion) for piece in board.pieces]
) )
if marker_id == self.MARKER_ON_BOARD_ID: if marker_id == self.MARKER_ON_BOARD_ID:
...@@ -310,6 +319,9 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor): ...@@ -310,6 +319,9 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor):
return board return board
def is_ready(self) -> bool:
return self.ready
def _board_from_corners(self) -> Board: def _board_from_corners(self) -> Board:
left_pieces = self._four_pieces_from_two_ends_vertical(self.bottom_left, self.top_left) left_pieces = self._four_pieces_from_two_ends_vertical(self.bottom_left, self.top_left)
right_pieces = self._four_pieces_from_two_ends_vertical(self.bottom_right, self.top_right) right_pieces = self._four_pieces_from_two_ends_vertical(self.bottom_right, self.top_right)
......
...@@ -71,20 +71,17 @@ class MemoryArucoFull(Memory): ...@@ -71,20 +71,17 @@ class MemoryArucoFull(Memory):
def __init__( def __init__(
self, self,
video_capture: cv.VideoCapture, video_capture: cv.VideoCapture,
piece_state_extractor: PieceStateExtractor = ArucoFullPieceStateExtractor(),
): ):
Memory.__init__(self, video_capture) Memory.__init__(self, video_capture)
self.piece_state_extractor = piece_state_extractor self.piece_state_extractor = ArucoFullPieceStateExtractor()
self.aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_1000) self.aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_1000)
self.aruco_params = cv.aruco.DetectorParameters_create() self.aruco_params = cv.aruco.DetectorParameters_create()
self.average_trigger = PieceTakenTrigger(self.piece_trigger) self.average_trigger = PieceTakenTrigger(self.piece_trigger)
self.frame_counter = 0
def process(self, img: np.ndarray) -> None: def process(self, img: np.ndarray) -> None:
# For the 5 first frames, store the corners # For the 5 first frames, store the corners
if not self.piece_state_extractor.ready: if not self.piece_state_extractor.ready:
self.piece_state_extractor.train(img) self.piece_state_extractor.train(img)
self.frame_counter += 1
return return
board = self.piece_state_extractor.get_board_with_visibility(img) board = self.piece_state_extractor.get_board_with_visibility(img)
...@@ -100,7 +97,7 @@ class MemoryArucoFull(Memory): ...@@ -100,7 +97,7 @@ class MemoryArucoFull(Memory):
def reset(self) -> None: def reset(self) -> None:
"""Reset board position""" """Reset board position"""
print("postion reset") print("postion reset")
self.frame_counter = 0 self.piece_state_extractor = ArucoFullPieceStateExtractor()
class MemoryArucoHalf(Memory): class MemoryArucoHalf(Memory):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment