import os

# ======================
# 配置
# ======================
LABEL_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/detect/val"  # 改成你的标签文件夹路径
BACKUP = False                  # 是否备份原文件

def convert_labels_to_zero(label_dir):
    for filename in os.listdir(label_dir):
        if not filename.endswith(".txt"):
            continue

        file_path = os.path.join(label_dir, filename)

        with open(file_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        new_lines = []
        for line in lines:
            line = line.strip()
            if not line:
                continue

            parts = line.split()
            if len(parts) < 5:
                # 非标准 yolo 行，原样保留
                new_lines.append(line)
                continue

            # 强制类别改为 0
            parts[0] = "0"
            new_lines.append(" ".join(parts))

        # 备份
        if BACKUP:
            backup_path = file_path + ".bak"
            if not os.path.exists(backup_path):
                os.rename(file_path, backup_path)
            save_path = file_path
        else:
            save_path = file_path

        with open(save_path, "w", encoding="utf-8") as f:
            f.write("\n".join(new_lines) + "\n")

        print(f"[OK] 已处理: {filename}")

if __name__ == "__main__":
    convert_labels_to_zero(LABEL_DIR)