From 955fa24b440d50b926ea1f94313ca6219a94d351 Mon Sep 17 00:00:00 2001
From: Adrien Lescourt <adrien@lescourt.net>
Date: Tue, 7 Dec 2021 14:35:56 +0100
Subject: [PATCH] Few tweaks
---
memory_lib/detectors/piece_state_detectors.py | 16 ++++++++++++++--
memory_lib/memory.py | 7 ++-----
2 files changed, 16 insertions(+), 7 deletions(-)
diff --git a/memory_lib/detectors/piece_state_detectors.py b/memory_lib/detectors/piece_state_detectors.py
index a2855dd..5479bef 100644
--- a/memory_lib/detectors/piece_state_detectors.py
+++ b/memory_lib/detectors/piece_state_detectors.py
@@ -31,6 +31,10 @@ class PieceStateExtractor(ABC):
def train(self, img: np.ndarray) -> None:
...
+ @abstractmethod
+ def is_ready(self) -> bool:
+ ...
+
class LightnessPieceStateDetector(PieceStateDetector):
""" Primitive detector based on lightness of average pixels around the piece"""
@@ -265,7 +269,12 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor):
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())
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)
+ 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)
if len(all_markers) == 16:
@@ -287,7 +296,7 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor):
)
# DEBUG
- cv.aruco.drawDetectedMarkers(img, corners)
+ # cv.aruco.drawDetectedMarkers(img, corners)
if ids is not None:
all_markers = []
@@ -298,7 +307,7 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor):
all_markers.append((marker_ids[0], center))
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]
)
if marker_id == self.MARKER_ON_BOARD_ID:
@@ -310,6 +319,9 @@ class ArucoFullPieceStateExtractor(PieceStateExtractor):
return board
+ def is_ready(self) -> bool:
+ return self.ready
+
def _board_from_corners(self) -> Board:
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)
diff --git a/memory_lib/memory.py b/memory_lib/memory.py
index e83d152..3375a96 100644
--- a/memory_lib/memory.py
+++ b/memory_lib/memory.py
@@ -71,20 +71,17 @@ class MemoryArucoFull(Memory):
def __init__(
self,
video_capture: cv.VideoCapture,
- piece_state_extractor: PieceStateExtractor = ArucoFullPieceStateExtractor(),
):
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_params = cv.aruco.DetectorParameters_create()
self.average_trigger = PieceTakenTrigger(self.piece_trigger)
- self.frame_counter = 0
def process(self, img: np.ndarray) -> None:
# For the 5 first frames, store the corners
if not self.piece_state_extractor.ready:
self.piece_state_extractor.train(img)
- self.frame_counter += 1
return
board = self.piece_state_extractor.get_board_with_visibility(img)
@@ -100,7 +97,7 @@ class MemoryArucoFull(Memory):
def reset(self) -> None:
"""Reset board position"""
print("postion reset")
- self.frame_counter = 0
+ self.piece_state_extractor = ArucoFullPieceStateExtractor()
class MemoryArucoHalf(Memory):
--
GitLab