107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
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() |