Separate rising and falling edges

in PieceTakenEdgeDetection

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:

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: