Minor fixes

rpi/models/detection/board_manager.py

@@ -0,0 +1,74 @@
+import cv2
+import numpy as np
+from typing import Tuple, Any
+
+from numpy import ndarray
+
+class BoardManager:
+
+    def process_frame(self, prediction: object, image : np.ndarray, scale_size: tuple[int, int]) -> tuple[ndarray, ndarray] | None:
+        try :
+            mask = self.__get_mask(prediction)
+            contour = self.__get_largest_contour(mask)
+            corners = self.__approx_corners(contour)
+            scaled_corners = self.__scale_corners(corners, mask.shape, image.shape)
+            ordered_corners = self.__order_corners(scaled_corners)
+            transformation_matrix = self.__calculte_transformation_matrix(ordered_corners, scale_size)
+            warped_corners = cv2.perspectiveTransform(
+                np.array(ordered_corners, np.float32).reshape(-1, 1, 2),
+                transformation_matrix
+            ).reshape(-1, 2)
+            return warped_corners, transformation_matrix
+        except Exception as e:
+            print(e)
+            return None
+
+    def __calculte_transformation_matrix(self, corners: np.ndarray, output_size : tuple[int, int]) -> np.ndarray:
+        width = output_size[0]
+        height = output_size[1]
+
+        dst = np.array([
+            [0, 0],  # top-left
+            [width - 1, 0],  # top-right
+            [width - 1, height - 1],  # bottom-right
+            [0, height - 1]  # bottom-left
+        ], dtype=np.float32)
+        return cv2.getPerspectiveTransform(corners, dst)
+
+    def __get_mask(self, pred: object) -> Any:
+        if pred.masks is None:
+            raise ValueError("Board contour is not 4 corners")
+        mask = pred.masks.data[0].cpu().numpy()
+        mask = (mask * 255).astype(np.uint8)
+        return mask
+
+    def __get_largest_contour(self, mask: np.ndarray) -> np.ndarray:
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if not contours:
+            raise ValueError("No contours found")
+        return max(contours, key=cv2.contourArea)
+
+    def __approx_corners(self, contour: np.ndarray) -> np.ndarray:
+        epsilon = 0.02 * cv2.arcLength(contour, True)
+        approx = cv2.approxPolyDP(contour, epsilon, True)
+        if len(approx) != 4:
+            raise ValueError("Board contour is not 4 corners")
+        return approx.reshape(4, 2)
+
+    def __scale_corners(self, pts: np.ndarray, mask_shape: Tuple[int, int], image_shape: Tuple[int, int, int]) -> np.ndarray:
+        mask_h, mask_w = mask_shape
+        img_h, img_w = image_shape[:2]
+        scale_x = img_w / mask_w
+        scale_y = img_h / mask_h
+        scaled_pts = [(int(p[0] * scale_x), int(p[1] * scale_y)) for p in pts]
+        return np.array(scaled_pts, dtype=np.float32)
+
+    def __order_corners(self, pts: np.ndarray) -> np.ndarray:
+        rect = np.zeros((4, 2), dtype="float32")
+        s = pts.sum(axis=1)
+        rect[0] = pts[np.argmin(s)]  # top-left
+        rect[2] = pts[np.argmax(s)]  # bottom-right
+        diff = np.diff(pts, axis=1)
+        rect[1] = pts[np.argmin(diff)]  # top-right
+        rect[3] = pts[np.argmax(diff)]  # bottom-left
+        return rect