import os

labels_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/edges/misis2025_cv_chess.v1i.yolov11/valid/labels"

labels_to_be_removed = [0, 1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

for filename in os.listdir(labels_dir):
    if not filename.endswith(".txt"):
        continue

    txt_path = os.path.join(labels_dir, filename)
    new_lines = []

    with open(txt_path, "r") as f:
        lines = f.readlines()

    for line in lines:
        parts = line.strip().split()

        if len(parts) < 5:
            continue

        cls = int(parts[0])
        if cls in labels_to_be_removed:
            print(f"{parts} found in {filename}")
            continue
        cls = 0
        new_lines.append(" ".join([str(cls)] + parts[1:]))

    with open(txt_path, "w") as f:
        f.write("\n".join(new_lines))

print("All label files have been adjusted!")