批量修改YOLO格式的标注类别

在这里插入图片描述

1.解决的问题

假如你有一个YOLO格式的数据集，标注类别为0，1，2，3四个类别标签。如果你想删除标签1，只保留0，2，3类别的标注信息，或者想将标签0和标签1合并为标签1，只剩下标签1，2，3。等类似的标签修改问题

2.python代码实现

import os# 原始txt文件存放的路径
txt_folder = "/dataset/source-sj/train/labels"
# 修改后txt文件的保存路径
new_txt_folder = "/dataset/source-sj/train/labels-1"# 确保新的保存文件夹存在，不存在则创建
if not os.path.exists(new_txt_folder):os.makedirs(new_txt_folder)# 遍历txt文件夹中的所有文件
for txt_file in os.listdir(txt_folder):if txt_file.endswith(".txt"):original_file_path = os.path.join(txt_folder, txt_file)new_file_path = os.path.join(new_txt_folder, txt_file)# 打开并读取原始文件内容with open(original_file_path, "r") as file:lines = file.readlines()new_lines = []for line in lines:components = line.strip().split()class_id = int(components[0])#在此处进行标签修改# 替换类别1为0，类别2为1，类别3为2if class_id == 1:components[0] = "0"elif class_id == 2:components[0] = "1"elif class_id == 3:components[0] = "1"new_lines.append(" ".join(components))# 将修改后的内容保存到新的文件中，确保不添加多余的空行with open(new_file_path, "w") as file:file.write("\n".join(new_lines))print("转换并保存完成！")