83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
import cv2
|
|
import numpy as np
|
|
from typing import Tuple
|
|
|
|
class BoardManager:
|
|
image: np.ndarray
|
|
|
|
def __init__(self, image : np.ndarray) -> None:
|
|
self.image = image
|
|
|
|
def extract_corners(self, prediction: object, scale_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Extrait le plateau warpé à partir de la prédiction et de l'image de base"""
|
|
|
|
mask = self.__get_mask(prediction)
|
|
contour = self.__get_largest_contour(mask)
|
|
corners = self.__approx_corners(contour)
|
|
scaled_corners = self.__scale_corners(corners, mask.shape, self.image.shape)
|
|
ordered_corners = self.__order_corners(scaled_corners)
|
|
transformation_matrix = self.__calculte_transformation_matrix(ordered_corners, scale_size)
|
|
return ordered_corners, transformation_matrix
|
|
|
|
def __calculte_transformation_matrix(self, corners: np.ndarray, output_size : tuple[int, int]) -> np.ndarray:
|
|
"""Calcule la matrice de perspective"""
|
|
|
|
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) -> np.ndarray:
|
|
"""Retourne le masque du plateau en uint8"""
|
|
|
|
if pred.masks is None:
|
|
raise ValueError("No board mask detected")
|
|
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:
|
|
"""Trouve le plus grand contour dans le masque"""
|
|
|
|
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:
|
|
"""Approxime les coins du plateau à 4 points"""
|
|
|
|
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:
|
|
"""Remappe les coins de la taille du masque vers la taille originale"""
|
|
|
|
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:
|
|
"""Ordonne les coins top-left, top-right, bottom-right, bottom-left"""
|
|
|
|
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
|