diff --git a/memory_lib/board.py b/memory_lib/board.py index 74a1ab9775cb2170c65c1390960f50eff79f72f5..b44294b70aff05c9f0e50336e7a1f35de9d49045 100644 --- a/memory_lib/board.py +++ b/memory_lib/board.py @@ -426,22 +426,24 @@ class ArucoPieceVisibilityDetector: ) -class AveragePieceTakenTrigger: +class PieceTakenDetectionEdgesTrigger: """Call 'trigger' function when after a piece is taken - It waits for 'avg' times the same piece detection state before triggering + It waits for 'rising_count' times the same piece detection state before triggering when we take a piece + It waits for 'falling_count' times the same piece detection state before triggering when we put back a piece TODO: replace this average impl with a counter to increase perf """ - def __init__(self, trigger: Callable[[int], None], avg: int = 5): + def __init__(self, trigger: Callable[[int], None], rising_count: int = 5, falling_count: int = 5): self.trigger = trigger - self.avg = avg + self.rising_count = rising_count + self.falling_count = falling_count self.last_boards: List[Board] = [] self.out_board: Optional[Board] = None def add_board(self, board: Board) -> None: - if len(self.last_boards) > self.avg * 2: # FIFO size + if len(self.last_boards) > max(self.rising_count, self.falling_count) + 1: # FIFO size self.last_boards.pop(0) self.last_boards.append(board) self.out_board = self._average_board() @@ -449,20 +451,20 @@ class AveragePieceTakenTrigger: def _average_board(self) -> Board: next_out_board = Board.from_board(self.last_boards[-1]) for idx in range(16): - avg = 0 + current_count = 0 for board in self.last_boards: - avg += int(board.pieces[idx].is_visible) + current_count += int(board.pieces[idx].is_visible) if self.out_board: if self.out_board.pieces[idx].is_visible: - if avg < 1: + if current_count < len(self.last_boards) - self.falling_count: next_out_board.pieces[idx].is_visible = False else: next_out_board.pieces[idx].is_visible = self.out_board.pieces[ idx ].is_visible else: - if avg > self.avg: + if current_count > self.rising_count: next_out_board.pieces[idx].is_visible = True self.trigger(idx) else: diff --git a/memory_lib/memory.py b/memory_lib/memory.py index a5280fe511b3327ba44790e08c0fcaebd0413aa5..8d2c61c2ad93ac44cdd3e8c0dea83be006e63766 100644 --- a/memory_lib/memory.py +++ b/memory_lib/memory.py @@ -18,7 +18,7 @@ from .board import ( ArucoPieceVisibilityDetector, ABMExtractor, AverageColorAndMovementPieceVisibilityDetector, - AveragePieceTakenTrigger, + PieceTakenDetectionEdgesTrigger, ) from .model import Board, PiecesObserver, Rect, Point @@ -76,7 +76,7 @@ class MemoryAruco(Memory): self.visibility_detector = visibility_detector self.aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_250) self.aruco_params = cv.aruco.DetectorParameters_create() - self.average_trigger = AveragePieceTakenTrigger(self.piece_trigger, 5) + self.average_trigger = PieceTakenDetectionEdgesTrigger(self.piece_trigger, 10, 40) self.frame_counter = 0 def process(self, img: np.ndarray) -> None: @@ -118,7 +118,7 @@ class MemoryABM(Memory): self.abm_extractor = abm_extractor self.visibility_detector = visibility_detector self.frame_counter = 0 - self.average_trigger = AveragePieceTakenTrigger(self.piece_trigger, 5) + self.average_trigger = PieceTakenDetectionEdgesTrigger(self.piece_trigger, 5) self.last_board = None def process(self, img: np.ndarray) -> None: