Add conversion to FEN

rpi/board-detector/board_manager.py

@@ -0,0 +1,82 @@
+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

rpi/board-detector/debug/realtime_detect.py

@@ -3,8 +3,8 @@ import cv2
 
 if __name__ == "__main__":
 
-    corner_model_path = "../assets/models/corner.pt"
-    pieces_model_path = "../assets/models/unified-nano-refined.pt"
+    corner_model_path = "../../assets/models/edges.pt"
+    pieces_model_path = "../../assets/models/unified-nano-refined.pt"
 
     print("Initializing model...")
     corner_model = YOLO(corner_model_path)

rpi/board-detector/detector.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+import os
+import random
+
+from ultralytics import YOLO
+import cv2
+
+class Detector :
+
+    def __init__(self, model_path):
+        self.model = YOLO(model_path)
+        self.used_height = 640
+        self.used_width = 640
+
+    def make_prediction(self, image_path):
+        return self.model.predict(source=image_path, conf=0.6)
+
+if __name__ == "__main__":
+    corner_model_path = "../assets/models/edges.pt"
+    pieces_model_path = "../assets/models/unified-nano-refined.pt"
+
+    corner_model = YOLO(corner_model_path)
+    pieces_model = YOLO(pieces_model_path)
+
+    img_folder = "../training/datasets/pieces/unified/test/images/"
+    save_folder = "./results"
+    os.makedirs(save_folder, exist_ok=True)
+
+    test_images = os.listdir(img_folder)
+
+    for i in range(0, 10):
+        rnd = random.randint(0, len(test_images) - 1)
+        img_path = os.path.join(img_folder, test_images[rnd])
+        save_path = os.path.join(save_folder, test_images[rnd])
+
+        image = cv2.imread(img_path)
+        height, width = image.shape[:2]
+
+        #fen = prediction_to_fen(results, height, width)
+        #print("Predicted FEN:", fen)
+
+        corner_result = corner_model.predict(source=image, conf=0.6)
+        pieces_result = pieces_model.predict(source=image, conf=0.6)
+
+        corner_annotated_image = corner_result[0].plot()
+        pieces_annotated_image = pieces_result[0].plot(img=corner_annotated_image)
+
+        cv2.imwrite(save_path, pieces_annotated_image)
+        #cv2.namedWindow("YOLO Predictions", cv2.WINDOW_NORMAL)
+        #cv2.imshow("YOLO Predictions", annotated_image)
+
+        cv2.waitKey(0)
+        cv2.destroyAllWindows()

rpi/board-detector/main.py

@@ -1,93 +1,107 @@
-#!/usr/bin/env python3
-import os
-import random
+import cv2
+import numpy as np
+from detector import Detector
+from board_manager import BoardManager
 
-from ultralytics import YOLO
+# -------------------- Pièces --------------------
+def extract_pieces(pieces_pred):
+    """Extrait les pièces avec leur bbox, sans remapping inutile"""
+    result = pieces_pred[0]
+    detections = []
+
+    for box in result.boxes:
+        # xywh en pixels de l'image originale
+        x, y, w, h = box.xywh[0].cpu().numpy()
+        label = result.names[int(box.cls[0])]
+        detections.append({"label": label, "bbox": (int(x), int(y), int(w), int(h))})
+
+    return detections
+
+import numpy as np
 import cv2
 
-# Map class names to FEN characters
-class_to_fen = {
-    'w_pawn': 'P',
-    'w_knight': 'N',
-    'w_bishop': 'B',
-    'w_rook': 'R',
-    'w_queen': 'Q',
-    'w_king': 'K',
-    'b_pawn': 'p',
-    'b_knight': 'n',
-    'b_bishop': 'b',
-    'b_rook': 'r',
-    'b_queen': 'q',
-    'b_king': 'k',
-}
+def pieces_to_board(detected_boxes, matrix, board_size=800):
+    board_array = [[None for _ in range(8)] for _ in range(8)]
 
-def prediction_to_fen(results, width, height):
+    for d in detected_boxes:
+        x, y, w, h = d["bbox"]
 
-    # Initialize empty board
-    board = [['' for _ in range(8)] for _ in range(8)]
+        # Points multiples sur la pièce pour stabilité
+        points = np.array([
+            [x + w/2, y + h*0.2],   # haut
+            [x + w/2, y + h/2],     # centre
+            [x + w/2, y + h*0.8]    # bas
+        ], dtype=np.float32).reshape(-1,1,2)
 
-    # Iterate through predictions
-    for result in results:
-        for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
-            x1, y1, x2, y2 = box.tolist()
-            class_name = model.names[int(cls)]
-            fen_char = class_to_fen.get(class_name)
+        # Transformation perspective
+        warped_points = cv2.perspectiveTransform(points, matrix)
+        wy_values = warped_points[:,0,1]  # coordonnées y après warp
 
-            if fen_char:
-                # Compute board square
-                col = int((x1 + x2) / 2 / (width / 8))
-                row = 7 - int((y1 + y2) / 2 / (height / 8))
-                board[row][col] = fen_char
-                print(f"[{class_name}] {fen_char} {row} {col}")
+        # Prendre le percentile haut (25%) pour éviter décalage
+        wy_percentile = np.percentile(wy_values, 25)
 
-    # Convert board to FEN
-    fen_rows = []
-    for row in board:
-        fen_row = ''
-        empty_count = 0
-        for square in row:
-            if square == '':
-                empty_count += 1
+        # Normaliser et calculer rank/file
+        nx = np.clip(np.mean(warped_points[:,0,0]) / board_size, 0, 0.999)
+        ny = np.clip(wy_percentile / board_size, 0, 0.999)
+
+        file = min(max(int(nx * 8), 0), 7)
+        rank = min(max(int(ny * 8), 0), 7)
+
+        board_array[rank][file] = d["label"]
+
+    return board_array
+
+
+def board_to_fen(board):
+    map_fen = {
+        "w_pawn": "P", "w_knight": "N", "w_bishop": "B",
+        "w_rook": "R", "w_queen": "Q", "w_king": "K",
+        "b_pawn": "p", "b_knight": "n", "b_bishop": "b",
+        "b_rook": "r", "b_queen": "q", "b_king": "k",
+    }
+    rows = []
+    for rank in board:
+        empty = 0
+        row = ""
+        for sq in rank:
+            if sq is None:
+                empty += 1
             else:
-                if empty_count > 0:
-                    fen_row += str(empty_count)
-                    empty_count = 0
-                fen_row += square
-        if empty_count > 0:
-            fen_row += str(empty_count)
-        fen_rows.append(fen_row)
-
-    # Join rows into a FEN string (default: white to move, all castling rights, no en passant)
-    fen_string = '/'.join(fen_rows) + ' w KQkq - 0 1'
-    return fen_string
-
+                if empty:
+                    row += str(empty)
+                    empty = 0
+                row += map_fen[sq]
+        if empty:
+            row += str(empty)
+        rows.append(row)
+    return "/".join(rows)
 
 if __name__ == "__main__":
-    model_path = "../assets/models/unified-nano-refined.pt"
-    img_folder = "../training/datasets/pieces/unified/test/images/"
-    save_folder = "./results"
-    os.makedirs(save_folder, exist_ok=True)
+    edges_detector = Detector("../assets/models/edges.pt")
+    pieces_detector = Detector("../assets/models/unified-nano-refined.pt")
+    #image_path = "./test/1.png"
+    image_path = "../training/datasets/pieces/unified/test/images/659_jpg.rf.0009cadea8df487a76d6960a28b9d811.jpg"
+    image = cv2.imread(image_path)
 
-    test_images = os.listdir(img_folder)
+    edges_pred = edges_detector.make_prediction(image_path)
+    pieces_pred = pieces_detector.make_prediction(image_path)
 
-    for i in range(0, 10):
-        rnd = random.randint(0, len(test_images) - 1)
-        img_path = os.path.join(img_folder, test_images[rnd])
-        save_path = os.path.join(save_folder, test_images[rnd])
+    remap_width = 800
+    remap_height = 800
 
-        img = cv2.imread(img_path)
-        height, width = img.shape[:2]
+    board_manager = BoardManager(image)
+    corners, matrix = board_manager.extract_corners(edges_pred[0], (remap_width, remap_height))
 
-        model = YOLO(model_path)
-        results = model.predict(source=img_path, conf=0.5)
+    detections = extract_pieces(pieces_pred)
 
-        #fen = prediction_to_fen(results, height, width)
-        #print("Predicted FEN:", fen)
+    board = pieces_to_board(detections, matrix, remap_width)
 
-        annotated_image = results[0].plot()
-        cv2.imwrite(save_path, annotated_image)
-        #cv2.namedWindow("YOLO Predictions", cv2.WINDOW_NORMAL)
-        #cv2.imshow("YOLO Predictions", annotated_image)
+    # FEN
+    fen = board_to_fen(board)
+    print("FEN:", fen)
 
-        cv2.waitKey(0)
-        cv2.destroyAllWindows()
+    frame = pieces_pred[0].plot()
+    cv2.namedWindow("Pred", cv2.WINDOW_NORMAL)
+    cv2.imshow("Pred", frame)
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()

rpi/training/training_v11_n.py

@@ -1,17 +1,17 @@
 from ultralytics import YOLO
 
 def main():
-    model = YOLO("models/unified-nano.pt")
+    model = YOLO("models/yolo11n-seg.pt")
     model.train(
-        data="./datasets/pieces/unified/data.yaml",
+        data="./datasets/edges/data.yaml",
         epochs=150,
         patience=20,
         imgsz=640,
-        batch=18,
+        batch=12,
         save_period=10,
         project="result",
-        name="unified-nano-refined",
-        exist_ok=True,
+        name="edges-nano",
+        exist_ok=False,
         device = 0
     )
 

rpi/training/utils/decrease_labels.py

@@ -1,13 +1,7 @@
 import os
 
-# --------------------------
-# Configuration
-# --------------------------
-labels_dir = "datasets/visiope/test/labels"
+labels_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/edges/chess board detection 2.v2i.yolov11"
 
-# --------------------------
-# Process each label file
-# --------------------------
 for filename in os.listdir(labels_dir):
     if not filename.endswith(".txt"):
         continue

rpi/training/utils/remove_labels.py

@@ -1,7 +1,8 @@
 import os
 
-labels_dir = "../datasets/corners/Outer Chess Corners.v1i.yolov11/valid/labels"
-label_to_be_removed = 1
+labels_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/edges/misis2025_cv_chess.v1i.yolov11/valid/labels"
+
+labels_to_be_removed = [0, 1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
 
 for filename in os.listdir(labels_dir):
     if not filename.endswith(".txt"):
@@ -20,13 +21,12 @@ for filename in os.listdir(labels_dir):
             continue
 
         cls = int(parts[0])
-        if cls == label_to_be_removed:
+        if cls in labels_to_be_removed:
             print(f"{parts} found in {filename}")
             continue
-
+        cls = 0
         new_lines.append(" ".join([str(cls)] + parts[1:]))
 
-    # Overwrite file with updated indices
     with open(txt_path, "w") as f:
         f.write("\n".join(new_lines))
 

rpi/training/utils/sort_labels.py

@@ -2,15 +2,12 @@ import os
 import shutil
 
 
-def copy_images(src, dest):
-    image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".gif"]
+def copy_files(src, dest):
 
     for filename in os.listdir(src):
-        if any(filename.lower().endswith(ext) for ext in image_extensions):
-            src_path = os.path.join(src, filename)
-            dst_path = os.path.join(dest, filename)
-            shutil.copy2(src_path, dst_path)
-
+        src_path = os.path.join(src, filename)
+        dst_path = os.path.join(dest, filename)
+        shutil.copy2(src_path, dst_path)
 
 def remap_labels(src, dest):
     count = 0
@@ -41,8 +38,8 @@ def remap_labels(src, dest):
 
 if __name__ == "__main__":
 
-    src_dir = "../datasets/pieces/visualizan/"
-    dest_dir = "../datasets/pieces/unified/"
+    src_dir = "../datasets/edges/misis2025_cv_chess.v1i.yolov11/"
+    dest_dir = "../datasets/edges/"
 
     reference_classes = [
         'w_pawn', 'w_knight', 'w_bishop', 'w_rook', 'w_queen', 'w_king',
@@ -67,5 +64,6 @@ if __name__ == "__main__":
         dst_image_folder = os.path.normpath(os.path.join(dest_full_path, "images"))
         dst_labels_folder = os.path.normpath(os.path.join(dest_full_path, "labels"))
 
-        copy_images(src_image_folder, dst_image_folder)
-        remap_labels(src_labels_folder, dst_labels_folder)
+        copy_files(src_image_folder, dst_image_folder)
+        copy_files(src_labels_folder, dst_labels_folder)
+        #remap_labels(src_labels_folder, dst_labels_folder)