diff --git a/rpi/assets/models/edges.pt b/rpi/assets/models/edges.pt new file mode 100644 index 00000000..3a2f0d89 Binary files /dev/null and b/rpi/assets/models/edges.pt differ diff --git a/rpi/assets/models/unified-nano-refined.pt b/rpi/assets/models/unified-nano-refined.pt index 749136d5..6fe4bfa0 100644 Binary files a/rpi/assets/models/unified-nano-refined.pt and b/rpi/assets/models/unified-nano-refined.pt differ diff --git a/rpi/board-detector/board_manager.py b/rpi/board-detector/board_manager.py new file mode 100644 index 00000000..4cf0f1f3 --- /dev/null +++ b/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 diff --git a/rpi/board-detector/debug/__init__.py b/rpi/board-detector/debug/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/rpi/board-detector/debug.py b/rpi/board-detector/debug/debug.py similarity index 100% rename from rpi/board-detector/debug.py rename to rpi/board-detector/debug/debug.py diff --git a/rpi/board-detector/realtime_detect.py b/rpi/board-detector/debug/realtime_detect.py similarity index 91% rename from rpi/board-detector/realtime_detect.py rename to rpi/board-detector/debug/realtime_detect.py index 39f26b2f..8e6f7982 100644 --- a/rpi/board-detector/realtime_detect.py +++ b/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) diff --git a/rpi/board-detector/detector.py b/rpi/board-detector/detector.py new file mode 100644 index 00000000..ea4a119b --- /dev/null +++ b/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() \ No newline at end of file diff --git a/rpi/board-detector/main.py b/rpi/board-detector/main.py index 535013c4..06531260 100644 --- a/rpi/board-detector/main.py +++ b/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() \ No newline at end of file + frame = pieces_pred[0].plot() + cv2.namedWindow("Pred", cv2.WINDOW_NORMAL) + cv2.imshow("Pred", frame) + cv2.waitKey(0) + cv2.destroyAllWindows() \ No newline at end of file diff --git a/rpi/training/training_v11_n.py b/rpi/training/training_v11_n.py index 535e9f59..6fca8a9a 100644 --- a/rpi/training/training_v11_n.py +++ b/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 ) diff --git a/rpi/training/utils/decrease_labels.py b/rpi/training/utils/decrease_labels.py index 3276ddb8..05368e06 100644 --- a/rpi/training/utils/decrease_labels.py +++ b/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 diff --git a/rpi/training/utils/remove_labels.py b/rpi/training/utils/remove_labels.py index 280e227c..babb78c1 100644 --- a/rpi/training/utils/remove_labels.py +++ b/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)) diff --git a/rpi/training/utils/sort_labels.py b/rpi/training/utils/sort_labels.py index 681c5035..87e19342 100644 --- a/rpi/training/utils/sort_labels.py +++ b/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)