Minor fixes
This commit is contained in:
@@ -1,82 +0,0 @@
|
||||
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
|
||||
@@ -1,46 +0,0 @@
|
||||
from ultralytics import YOLO
|
||||
import cv2
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
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)
|
||||
pieces_model = YOLO(pieces_model_path)
|
||||
|
||||
print("Initializing camera...")
|
||||
cap = cv2.VideoCapture(0)
|
||||
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
|
||||
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
|
||||
cap.set(cv2.CAP_PROP_FPS, 30)
|
||||
|
||||
print("Initialized")
|
||||
if not cap.isOpened():
|
||||
print("Error: Could not open camera")
|
||||
exit()
|
||||
|
||||
cv2.namedWindow("Predictions", cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
print("Error: Failed to grab frame")
|
||||
break
|
||||
|
||||
# Optional: resize frame to improve YOLO performance
|
||||
# frame = cv2.resize(frame, (416, 416))
|
||||
|
||||
corner_result = corner_model.predict(source=frame, conf=0.6)
|
||||
pieces_result = pieces_model.predict(source=frame, conf=0.6)
|
||||
|
||||
corner_annotated_frame = corner_result[0].plot()
|
||||
pieces_annotated_frame = pieces_result[0].plot(img=corner_annotated_frame)
|
||||
|
||||
cv2.imshow("Predictions", pieces_annotated_frame)
|
||||
cv2.resizeWindow("Predictions", 640, 640)
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
@@ -1,53 +0,0 @@
|
||||
#!/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,107 +0,0 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from detector import Detector
|
||||
from board_manager import BoardManager
|
||||
|
||||
# -------------------- 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
|
||||
|
||||
def pieces_to_board(detected_boxes, matrix, board_size=800):
|
||||
board_array = [[None for _ in range(8)] for _ in range(8)]
|
||||
|
||||
for d in detected_boxes:
|
||||
x, y, w, h = d["bbox"]
|
||||
|
||||
# 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)
|
||||
|
||||
# Transformation perspective
|
||||
warped_points = cv2.perspectiveTransform(points, matrix)
|
||||
wy_values = warped_points[:,0,1] # coordonnées y après warp
|
||||
|
||||
# Prendre le percentile haut (25%) pour éviter décalage
|
||||
wy_percentile = np.percentile(wy_values, 25)
|
||||
|
||||
# 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:
|
||||
row += str(empty)
|
||||
empty = 0
|
||||
row += map_fen[sq]
|
||||
if empty:
|
||||
row += str(empty)
|
||||
rows.append(row)
|
||||
return "/".join(rows)
|
||||
|
||||
if __name__ == "__main__":
|
||||
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)
|
||||
|
||||
edges_pred = edges_detector.make_prediction(image_path)
|
||||
pieces_pred = pieces_detector.make_prediction(image_path)
|
||||
|
||||
remap_width = 800
|
||||
remap_height = 800
|
||||
|
||||
board_manager = BoardManager(image)
|
||||
corners, matrix = board_manager.extract_corners(edges_pred[0], (remap_width, remap_height))
|
||||
|
||||
detections = extract_pieces(pieces_pred)
|
||||
|
||||
board = pieces_to_board(detections, matrix, remap_width)
|
||||
|
||||
# FEN
|
||||
fen = board_to_fen(board)
|
||||
print("FEN:", fen)
|
||||
|
||||
frame = pieces_pred[0].plot()
|
||||
cv2.namedWindow("Pred", cv2.WINDOW_NORMAL)
|
||||
cv2.imshow("Pred", frame)
|
||||
cv2.waitKey(0)
|
||||
cv2.destroyAllWindows()
|
||||
89
rpi/board-detector/realtime_detect.py
Normal file
89
rpi/board-detector/realtime_detect.py
Normal file
@@ -0,0 +1,89 @@
|
||||
import time
|
||||
|
||||
from models.detection.detector import *
|
||||
from models.detection.board_manager import *
|
||||
from models.detection.pieces_manager import *
|
||||
|
||||
def print_board_console(board_array):
|
||||
files = "abcdefgh"
|
||||
print(" " + " ".join(files))
|
||||
print(" +-----------------+")
|
||||
|
||||
for r in range(8):
|
||||
row_str = str(8 - r) + "|"
|
||||
for f in range(8):
|
||||
piece = board_array[r][f]
|
||||
if piece is None:
|
||||
row_str += ". "
|
||||
else:
|
||||
if piece[0] == "w":
|
||||
piece = piece.upper()
|
||||
row_str += piece[2] + " "
|
||||
row_str += "|"
|
||||
print(row_str)
|
||||
print(" +-----------------+")
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
print("Initializing models...")
|
||||
edges_detector = Detector("../assets/models/edges.pt")
|
||||
pieces_detector = Detector("../assets/models/unified-nano-refined.pt")
|
||||
|
||||
pieces_manager = PiecesManager()
|
||||
board_manager = BoardManager()
|
||||
|
||||
print("Initializing camera...")
|
||||
cap = cv2.VideoCapture(0)
|
||||
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
|
||||
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
|
||||
cap.set(cv2.CAP_PROP_FPS, 30)
|
||||
|
||||
print("Initialized")
|
||||
if not cap.isOpened():
|
||||
print("Error: Could not open camera")
|
||||
exit()
|
||||
|
||||
cv2.namedWindow("Predictions", cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
time.sleep(1)
|
||||
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
print("Error: Failed to grab frame")
|
||||
break
|
||||
|
||||
edges_pred = edges_detector.make_prediction(frame)
|
||||
pieces_pred = pieces_detector.make_prediction(frame)
|
||||
|
||||
edges_annotated_frame = edges_pred[0].plot()
|
||||
pieces_annotated_frame = pieces_pred[0].plot(img=edges_annotated_frame)
|
||||
|
||||
cv2.imshow("Predictions", pieces_annotated_frame)
|
||||
cv2.resizeWindow("Predictions", 640, 640)
|
||||
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
remap_width = 800
|
||||
remap_height = 800
|
||||
|
||||
board_manager = BoardManager()
|
||||
result = board_manager.process_frame(edges_pred[0], frame, (remap_width, remap_height))
|
||||
if result is None:
|
||||
continue
|
||||
|
||||
corners = result[0]
|
||||
matrix = result[1]
|
||||
|
||||
detections = pieces_manager.extract_pieces(pieces_pred)
|
||||
|
||||
board = pieces_manager.pieces_to_board(detections, matrix, (remap_width, remap_height))
|
||||
|
||||
fen = pieces_manager.board_to_fen(board)
|
||||
print("FEN:", fen)
|
||||
|
||||
print_board_console(board)
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
Reference in New Issue
Block a user