Untested aruco full stuff

memory_lib/__init__.py

-from .memory import MemoryABM, MemoryAruco
+from .memory import MemoryABM, MemoryAruco, MemoryArucoFull

memory_lib/board.py

@@ -19,7 +19,7 @@ from .geometry import (
     get_segment_center,
     get_segement_point_at_dist,
 )
-from .model import Board, Point, LineEq, Piece, Rect
+from .model import Board, Point, LineEq, Piece, Rect, PieceState
 
 
 class BoardFactory(ABC):
@@ -332,8 +332,103 @@ class AverageColorAndMovementPieceVisibilityDetector(PieceVisibilityDetector):
         return False
 
 
+class ArucoFullPieceVisibilityDetector:
+    """ Piece detector based on Aruco markers above and below every piece
+
+    TODO:
+    for now this detector does not work as a PieceVisibilityDetector
+    This detector directly send a board with updated visibility
+    """
+
+    def __init__(self):
+        self.aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_250)
+        self.aruco_params = cv.aruco.DetectorParameters_create()
+        self.top_left: Point = Point(-1, -1)
+        self.top_right: Point = Point(-1, -1)
+        self.bottom_left: Point = Point(-1, -1)
+        self.bottom_right: Point = Point(-1, -1)
+        self.board = None
+
+    def train(self, img: np.ndarray) -> None:
+        """ Store the position of the 4 markers below the the 4 pieces in the corners"""
+        corners, ids, rejected = cv.aruco.detectMarkers(
+            img, self.aruco_dict, parameters=self.aruco_params
+        )
+
+        if ids is not None:
+            all_markers = []
+            for i in range(len(ids)):
+                r = ((int(x), int(y)) for x, y in corners[i][0].tolist())
+                rect = Rect.from_corners(*(Point.from_tuple(t) for t in r))
+                center = get_rect_center(rect)
+                all_markers.append(center)
+
+            # we have the 4 corners
+            if len(all_markers) >= 4:
+                vertical_sort = sorted(all_markers, key=lambda point: point.y)
+                top_markers = vertical_sort[:2]
+                bottom_markers = vertical_sort[2:]
+                self.top_left = min(top_markers, key=lambda point: point.x)
+                self.top_right = max(top_markers, key=lambda point: point.x)
+                self.bottom_left = min(bottom_markers, key=lambda point: point.x)
+                self.bottom_right = max(bottom_markers, key=lambda point: point.x)
+
+    def get_board_with_visibility(self, img: np.ndarray) -> Board:
+        """Compute a new board from the 4 corners stored during training, with all pieces visibility"""
+        board = self._board_from_corners()
+
+        corners, ids, rejected = cv.aruco.detectMarkers(
+            img, self.aruco_dict, parameters=self.aruco_params
+        )
+
+        if ids is not None:
+            all_markers = []
+            for i in range(len(ids)):
+                r = ((int(x), int(y)) for x, y in corners[i][0].tolist())
+                rect = Rect.from_corners(*(Point.from_tuple(t) for t in r))
+                center = get_rect_center(rect)
+                all_markers.append(center)
+
+            for marker in all_markers:
+                closest_piece_idx = np.argmin(
+                    [get_segment_size(marker, piece.postion) for piece in board.pieces]
+                )
+                board.pieces[closest_piece_idx].state = PieceState.OFF_BOARD
+
+        return board
+
+    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)
+        row_0 = self._four_pieces_from_two_ends_horizontal(left_pieces[0].postion, right_pieces[0].postion)
+        row_1 = self._four_pieces_from_two_ends_horizontal(left_pieces[1].postion, right_pieces[1].postion)
+        row_2 = self._four_pieces_from_two_ends_horizontal(left_pieces[2].postion, right_pieces[2].postion)
+        row_3 = self._four_pieces_from_two_ends_horizontal(left_pieces[3].postion, right_pieces[3].postion)
+        return Board(row_0 + row_1 + row_2 + row_3)
+
+    @staticmethod
+    def _four_pieces_from_two_ends_vertical(p1: Point, p2: Point) -> Tuple[Piece, ...]:
+        """Get 4 pieces from 2 ends pieces"""
+        return (
+            Piece(p1),
+            Piece(get_segement_point_at_dist(p1, p2, 0.45)),
+            Piece(get_segement_point_at_dist(p2, p1, 0.20)),
+            Piece(p2),
+        )
+
+    @staticmethod
+    def _four_pieces_from_two_ends_horizontal(p1: Point, p2: Point) -> Tuple[Piece, ...]:
+        """Get 4 pieces from 2 ends pieces"""
+        return (
+            Piece(p1),
+            Piece(get_segement_point_at_dist(p1, p2, 0.33)),
+            Piece(get_segement_point_at_dist(p2, p1, 0.33)),
+            Piece(p2),
+        )
+
+
 class ArucoPieceVisibilityDetector:
-    """ Piece detector based Aruco marker
+    """ Piece detector based on Aruco marker below every piece
 
     TODO:
     for now this detector does not work as a PieceVisibilityDetector
@@ -477,3 +572,56 @@ class PieceTakenDetectionEdgesTrigger:
             #     print([int(b.pieces[idx].is_visible) for b in self.last_boards])
 
         return next_out_board
+
+
+class PieceTakenDetectionEdgesTriggerTestOnArucoFull:
+    """Call 'trigger' function when after a piece is taken
+
+    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], rising_count: int = 5, falling_count: int = 5):
+        self.trigger = trigger
+        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) > 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()
+
+    def _average_board(self) -> Board:
+        next_out_board = Board.from_board(self.last_boards[-1])
+        for idx in range(16):
+            current_count = 0
+            for board in self.last_boards:
+                current_count += int(board.pieces[idx].state == PieceState.OFF_BOARD)
+
+            if self.out_board:
+                if self.out_board.pieces[idx].state == PieceState.OFF_BOARD:
+                    if current_count < len(self.last_boards) - self.falling_count:
+                        next_out_board.pieces[idx].state = PieceState.ON_BOARD
+                    else:
+                        next_out_board.pieces[idx].state = self.out_board.pieces[
+                            idx
+                        ].state
+                else:
+                    if current_count > self.rising_count:
+                        next_out_board.pieces[idx].state = PieceState.OFF_BOARD
+                        self.trigger(idx)
+                    else:
+                        next_out_board.pieces[idx].state = self.out_board.pieces[
+                            idx
+                        ].state
+
+            # if idx == 5:
+            #     print(avg, next_out_board.pieces[idx].is_visible, self.last_board.pieces[idx].is_visible)
+            #     print([int(b.pieces[idx].is_visible) for b in self.last_boards])
+
+        return next_out_board

memory_lib/memory.py

@@ -16,9 +16,10 @@ from .board import (
     LightnessPieceVisibilityDetector,
     AverageColorPieceVisibilityDetector,
     ArucoPieceVisibilityDetector,
+    ArucoFullPieceVisibilityDetector,
     ABMExtractor,
     AverageColorAndMovementPieceVisibilityDetector,
-    PieceTakenDetectionEdgesTrigger,
+    PieceTakenDetectionEdgesTrigger, PieceTakenDetectionEdgesTriggerTestOnArucoFull,
 )
 from .model import Board, PiecesObserver, Rect, Point
 
@@ -70,6 +71,46 @@ class Memory(ABC, threading.Thread, PiecesObserver):
         ...
 
 
+class MemoryArucoFull(Memory):
+    """Memory pieces detection based on Aruco above and below every piece"""
+
+    def __init__(
+            self,
+            video_capture: cv.VideoCapture,
+            visibility_detector: ArucoPieceVisibilityDetector = ArucoFullPieceVisibilityDetector(),
+    ):
+        Memory.__init__(self, video_capture)
+        self.visibility_detector = visibility_detector
+        self.aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_1000)
+        self.aruco_params = cv.aruco.DetectorParameters_create()
+        self.average_trigger = PieceTakenDetectionEdgesTriggerTestOnArucoFull(
+            self.piece_trigger, 10, 40
+        )
+        self.frame_counter = 0
+
+    def process(self, img: np.ndarray) -> None:
+        # For the 5 first frames, store the corners
+        if self.frame_counter < 10:
+            self.visibility_detector.train(img)
+            self.frame_counter += 1
+            return
+
+        board = self.visibility_detector.get_board_with_visibility(img)
+        # self.board_trigger(board, img)
+        self.average_trigger.add_board(board)
+
+        # TODO: debug
+        # cv.aruco.drawDetectedMarkers(img, corners)
+        self.average_trigger.out_board.draw(img)
+        self.board_trigger(board, img)
+        cv.imshow("marker", img)
+
+    def reset(self) -> None:
+        """Reset board position"""
+        print("postion reset")
+        self.frame_counter = 0
+
+
 class MemoryAruco(Memory):
     """Memory pieces detection based on Aruco below every piece"""
 

memory_lib/model.py

 from dataclasses import dataclass
+from enum import Enum
 from typing import Tuple, Callable
 import numpy as np
 import cv2 as cv
@@ -39,20 +40,32 @@ class LineEq:
     b: float
 
 
+class PieceState(Enum):
+    ON_BOARD = 1
+    OFF_BOARD = 2
+    UNKNOWN = 3
+
+
 @dataclass
 class Piece:
     postion: Point = Point(0, 0)
-    is_visible: bool = False
+    state: PieceState = PieceState.UNKNOWN
 
 
 @dataclass
 class Board:
     pieces: Tuple[Piece, ...]  # All 16 pieces
 
+    _draw_colors = {
+        PieceState.ON_BOARD: (255, 0, 0),
+        PieceState.OFF_BOARD: (0, 255, 0),
+        PieceState.UNKNOWN: (0, 0, 255),
+    }
+
     def draw(self, img: np.ndarray):
         """ Draw circle for each piece position. Green if the piece is visible, red otherwise"""
         for piece in self.pieces:
-            color = (0, 255, 0) if piece.is_visible else (0, 0, 255)
+            color = self._draw_colors[piece.state]
             cv.circle(img, piece.postion.to_tuple(), 10, color)
 
     def update_pieces(self, img: np.ndarray, piece_detector) -> None:
@@ -61,7 +74,7 @@ class Board:
 
     @classmethod
     def from_board(cls, board: "Board") -> "Board":
-        pieces = tuple(Piece(p.postion, p.is_visible) for p in board.pieces)
+        pieces = tuple(Piece(p.postion, p.state) for p in board.pieces)
         return cls(pieces)
 
 

memory_lib/tests/aruco.py

 import cv2 as cv
 import numpy as np
 
-# cap = cv.VideoCapture(0)
-cap = cv.VideoCapture("res/webcam_05_aruco.avi")
-aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_250)
+cap = cv.VideoCapture(0)
+# cap = cv.VideoCapture("res/webcam_05_aruco.avi")
+aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_1000)
 aruco_params = cv.aruco.DetectorParameters_create()
 
 while True:

memory_lib/tests/run_memory.py

@@ -35,11 +35,11 @@ def main():
     # memory = memory_lib.MemoryABM(cap)
     # cap = cv.VideoCapture("res/webcam_05_aruco.avi")
 
-    # cap = cv.VideoCapture(0)
-    cap = cv.VideoCapture("res/webcam_11_aruco.avi")
-    memory = memory_lib.MemoryAruco(cap)
+    cap = cv.VideoCapture(0)
+    # cap = cv.VideoCapture("res/webcam_11_aruco.avi")
+    memory = memory_lib.MemoryArucoFull(cap)
 
-    memory.bind(hello_piece)
+    memory.bind_pieces(hello_piece)
     memory.start()
 
     # process_vid("res/webcam_03_shadow.avi", memory)

pyproject.toml

+[tool.pyright]
+venvPath = "."
+venv = "venv"