Add conversion to FEN
This commit is contained in:
BIN
rpi/assets/models/edges.pt
Normal file
BIN
rpi/assets/models/edges.pt
Normal file
Binary file not shown.
Binary file not shown.
82
rpi/board-detector/board_manager.py
Normal file
82
rpi/board-detector/board_manager.py
Normal file
@@ -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
|
||||||
0
rpi/board-detector/debug/__init__.py
Normal file
0
rpi/board-detector/debug/__init__.py
Normal file
@@ -3,8 +3,8 @@ import cv2
|
|||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|
||||||
corner_model_path = "../assets/models/corner.pt"
|
corner_model_path = "../../assets/models/edges.pt"
|
||||||
pieces_model_path = "../assets/models/unified-nano-refined.pt"
|
pieces_model_path = "../../assets/models/unified-nano-refined.pt"
|
||||||
|
|
||||||
print("Initializing model...")
|
print("Initializing model...")
|
||||||
corner_model = YOLO(corner_model_path)
|
corner_model = YOLO(corner_model_path)
|
||||||
53
rpi/board-detector/detector.py
Normal file
53
rpi/board-detector/detector.py
Normal file
@@ -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()
|
||||||
@@ -1,93 +1,107 @@
|
|||||||
#!/usr/bin/env python3
|
import cv2
|
||||||
import os
|
import numpy as np
|
||||||
import random
|
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
|
import cv2
|
||||||
|
|
||||||
# Map class names to FEN characters
|
def pieces_to_board(detected_boxes, matrix, board_size=800):
|
||||||
class_to_fen = {
|
board_array = [[None for _ in range(8)] for _ in range(8)]
|
||||||
'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 prediction_to_fen(results, width, height):
|
for d in detected_boxes:
|
||||||
|
x, y, w, h = d["bbox"]
|
||||||
|
|
||||||
# Initialize empty board
|
# Points multiples sur la pièce pour stabilité
|
||||||
board = [['' for _ in range(8)] for _ in range(8)]
|
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
|
# Transformation perspective
|
||||||
for result in results:
|
warped_points = cv2.perspectiveTransform(points, matrix)
|
||||||
for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
|
wy_values = warped_points[:,0,1] # coordonnées y après warp
|
||||||
x1, y1, x2, y2 = box.tolist()
|
|
||||||
class_name = model.names[int(cls)]
|
|
||||||
fen_char = class_to_fen.get(class_name)
|
|
||||||
|
|
||||||
if fen_char:
|
# Prendre le percentile haut (25%) pour éviter décalage
|
||||||
# Compute board square
|
wy_percentile = np.percentile(wy_values, 25)
|
||||||
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}")
|
|
||||||
|
|
||||||
# Convert board to FEN
|
# Normaliser et calculer rank/file
|
||||||
fen_rows = []
|
nx = np.clip(np.mean(warped_points[:,0,0]) / board_size, 0, 0.999)
|
||||||
for row in board:
|
ny = np.clip(wy_percentile / board_size, 0, 0.999)
|
||||||
fen_row = ''
|
|
||||||
empty_count = 0
|
file = min(max(int(nx * 8), 0), 7)
|
||||||
for square in row:
|
rank = min(max(int(ny * 8), 0), 7)
|
||||||
if square == '':
|
|
||||||
empty_count += 1
|
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:
|
else:
|
||||||
if empty_count > 0:
|
if empty:
|
||||||
fen_row += str(empty_count)
|
row += str(empty)
|
||||||
empty_count = 0
|
empty = 0
|
||||||
fen_row += square
|
row += map_fen[sq]
|
||||||
if empty_count > 0:
|
if empty:
|
||||||
fen_row += str(empty_count)
|
row += str(empty)
|
||||||
fen_rows.append(fen_row)
|
rows.append(row)
|
||||||
|
return "/".join(rows)
|
||||||
# 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 __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
model_path = "../assets/models/unified-nano-refined.pt"
|
edges_detector = Detector("../assets/models/edges.pt")
|
||||||
img_folder = "../training/datasets/pieces/unified/test/images/"
|
pieces_detector = Detector("../assets/models/unified-nano-refined.pt")
|
||||||
save_folder = "./results"
|
#image_path = "./test/1.png"
|
||||||
os.makedirs(save_folder, exist_ok=True)
|
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):
|
remap_width = 800
|
||||||
rnd = random.randint(0, len(test_images) - 1)
|
remap_height = 800
|
||||||
img_path = os.path.join(img_folder, test_images[rnd])
|
|
||||||
save_path = os.path.join(save_folder, test_images[rnd])
|
|
||||||
|
|
||||||
img = cv2.imread(img_path)
|
board_manager = BoardManager(image)
|
||||||
height, width = img.shape[:2]
|
corners, matrix = board_manager.extract_corners(edges_pred[0], (remap_width, remap_height))
|
||||||
|
|
||||||
model = YOLO(model_path)
|
detections = extract_pieces(pieces_pred)
|
||||||
results = model.predict(source=img_path, conf=0.5)
|
|
||||||
|
|
||||||
#fen = prediction_to_fen(results, height, width)
|
board = pieces_to_board(detections, matrix, remap_width)
|
||||||
#print("Predicted FEN:", fen)
|
|
||||||
|
|
||||||
annotated_image = results[0].plot()
|
# FEN
|
||||||
cv2.imwrite(save_path, annotated_image)
|
fen = board_to_fen(board)
|
||||||
#cv2.namedWindow("YOLO Predictions", cv2.WINDOW_NORMAL)
|
print("FEN:", fen)
|
||||||
#cv2.imshow("YOLO Predictions", annotated_image)
|
|
||||||
|
|
||||||
cv2.waitKey(0)
|
frame = pieces_pred[0].plot()
|
||||||
cv2.destroyAllWindows()
|
cv2.namedWindow("Pred", cv2.WINDOW_NORMAL)
|
||||||
|
cv2.imshow("Pred", frame)
|
||||||
|
cv2.waitKey(0)
|
||||||
|
cv2.destroyAllWindows()
|
||||||
@@ -1,17 +1,17 @@
|
|||||||
from ultralytics import YOLO
|
from ultralytics import YOLO
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
model = YOLO("models/unified-nano.pt")
|
model = YOLO("models/yolo11n-seg.pt")
|
||||||
model.train(
|
model.train(
|
||||||
data="./datasets/pieces/unified/data.yaml",
|
data="./datasets/edges/data.yaml",
|
||||||
epochs=150,
|
epochs=150,
|
||||||
patience=20,
|
patience=20,
|
||||||
imgsz=640,
|
imgsz=640,
|
||||||
batch=18,
|
batch=12,
|
||||||
save_period=10,
|
save_period=10,
|
||||||
project="result",
|
project="result",
|
||||||
name="unified-nano-refined",
|
name="edges-nano",
|
||||||
exist_ok=True,
|
exist_ok=False,
|
||||||
device = 0
|
device = 0
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -1,13 +1,7 @@
|
|||||||
import os
|
import os
|
||||||
|
|
||||||
# --------------------------
|
labels_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/edges/chess board detection 2.v2i.yolov11"
|
||||||
# Configuration
|
|
||||||
# --------------------------
|
|
||||||
labels_dir = "datasets/visiope/test/labels"
|
|
||||||
|
|
||||||
# --------------------------
|
|
||||||
# Process each label file
|
|
||||||
# --------------------------
|
|
||||||
for filename in os.listdir(labels_dir):
|
for filename in os.listdir(labels_dir):
|
||||||
if not filename.endswith(".txt"):
|
if not filename.endswith(".txt"):
|
||||||
continue
|
continue
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
import os
|
import os
|
||||||
|
|
||||||
labels_dir = "../datasets/corners/Outer Chess Corners.v1i.yolov11/valid/labels"
|
labels_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/edges/misis2025_cv_chess.v1i.yolov11/valid/labels"
|
||||||
label_to_be_removed = 1
|
|
||||||
|
labels_to_be_removed = [0, 1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
|
||||||
|
|
||||||
for filename in os.listdir(labels_dir):
|
for filename in os.listdir(labels_dir):
|
||||||
if not filename.endswith(".txt"):
|
if not filename.endswith(".txt"):
|
||||||
@@ -20,13 +21,12 @@ for filename in os.listdir(labels_dir):
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
cls = int(parts[0])
|
cls = int(parts[0])
|
||||||
if cls == label_to_be_removed:
|
if cls in labels_to_be_removed:
|
||||||
print(f"{parts} found in {filename}")
|
print(f"{parts} found in {filename}")
|
||||||
continue
|
continue
|
||||||
|
cls = 0
|
||||||
new_lines.append(" ".join([str(cls)] + parts[1:]))
|
new_lines.append(" ".join([str(cls)] + parts[1:]))
|
||||||
|
|
||||||
# Overwrite file with updated indices
|
|
||||||
with open(txt_path, "w") as f:
|
with open(txt_path, "w") as f:
|
||||||
f.write("\n".join(new_lines))
|
f.write("\n".join(new_lines))
|
||||||
|
|
||||||
|
|||||||
@@ -2,15 +2,12 @@ import os
|
|||||||
import shutil
|
import shutil
|
||||||
|
|
||||||
|
|
||||||
def copy_images(src, dest):
|
def copy_files(src, dest):
|
||||||
image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".gif"]
|
|
||||||
|
|
||||||
for filename in os.listdir(src):
|
for filename in os.listdir(src):
|
||||||
if any(filename.lower().endswith(ext) for ext in image_extensions):
|
src_path = os.path.join(src, filename)
|
||||||
src_path = os.path.join(src, filename)
|
dst_path = os.path.join(dest, filename)
|
||||||
dst_path = os.path.join(dest, filename)
|
shutil.copy2(src_path, dst_path)
|
||||||
shutil.copy2(src_path, dst_path)
|
|
||||||
|
|
||||||
|
|
||||||
def remap_labels(src, dest):
|
def remap_labels(src, dest):
|
||||||
count = 0
|
count = 0
|
||||||
@@ -41,8 +38,8 @@ def remap_labels(src, dest):
|
|||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|
||||||
src_dir = "../datasets/pieces/visualizan/"
|
src_dir = "../datasets/edges/misis2025_cv_chess.v1i.yolov11/"
|
||||||
dest_dir = "../datasets/pieces/unified/"
|
dest_dir = "../datasets/edges/"
|
||||||
|
|
||||||
reference_classes = [
|
reference_classes = [
|
||||||
'w_pawn', 'w_knight', 'w_bishop', 'w_rook', 'w_queen', 'w_king',
|
'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_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"))
|
dst_labels_folder = os.path.normpath(os.path.join(dest_full_path, "labels"))
|
||||||
|
|
||||||
copy_images(src_image_folder, dst_image_folder)
|
copy_files(src_image_folder, dst_image_folder)
|
||||||
remap_labels(src_labels_folder, dst_labels_folder)
|
copy_files(src_labels_folder, dst_labels_folder)
|
||||||
|
#remap_labels(src_labels_folder, dst_labels_folder)
|
||||||
|
|||||||
Reference in New Issue
Block a user