71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
import os
|
|
import cv2
|
|
from ultralytics import YOLO
|
|
|
|
# --------------------------
|
|
# Configuration
|
|
# --------------------------
|
|
model_path = "models/bck/best-3.pt" # your trained YOLO model
|
|
images_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/universe/train/images"
|
|
labels_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/universe/train/labels"
|
|
img_width = 640
|
|
img_height = 640
|
|
|
|
os.makedirs(labels_dir, exist_ok=True)
|
|
|
|
# --------------------------
|
|
# Load model
|
|
# --------------------------
|
|
model = YOLO(model_path)
|
|
|
|
# --------------------------
|
|
# Mapping YOLO class index -> piece name (optional)
|
|
# --------------------------
|
|
names = ['w_pawn','w_knight','w_bishop','w_rook','w_queen','w_king',
|
|
'b_pawn','b_knight','b_bishop','b_rook','b_queen','b_king']
|
|
|
|
# --------------------------
|
|
# Process images
|
|
# --------------------------
|
|
for img_file in os.listdir(images_dir):
|
|
if not img_file.lower().endswith((".png", ".jpg", ".jpeg")):
|
|
continue
|
|
|
|
img_path = os.path.join(images_dir, img_file)
|
|
img = cv2.imread(img_path)
|
|
if img is None:
|
|
print(f"Failed to read {img_file}")
|
|
continue
|
|
|
|
height, width = img.shape[:2]
|
|
|
|
# Run YOLO detection
|
|
results = model(img)
|
|
res = results[0]
|
|
|
|
lines = []
|
|
boxes = res.boxes.xyxy.cpu().numpy() # [x1, y1, x2, y2]
|
|
classes = res.boxes.cls.cpu().numpy()
|
|
confs = res.boxes.conf.cpu().numpy()
|
|
|
|
for box, cls, conf in zip(boxes, classes, confs):
|
|
if conf < 0.5: # skip low-confidence predictions
|
|
continue
|
|
|
|
x1, y1, x2, y2 = box
|
|
x_center = (x1 + x2) / 2 / width
|
|
y_center = (y1 + y2) / 2 / height
|
|
w_norm = (x2 - x1) / width
|
|
h_norm = (y2 - y1) / height
|
|
|
|
lines.append(f"{int(cls)} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}")
|
|
|
|
# Save YOLO .txt file with same basename as image
|
|
txt_path = os.path.join(labels_dir, os.path.splitext(img_file)[0] + ".txt")
|
|
with open(txt_path, "w") as f:
|
|
f.write("\n".join(lines))
|
|
|
|
print(f"Pre-labeled {img_file} -> {txt_path}")
|
|
|
|
print("All images have been pre-labeled!")
|