Idk

rpi/training/labelizer.py

@@ -0,0 +1,70 @@
+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!")