批量删除并更改yolo格式的标签类别

在用yolo做目标检测项目时，我们经常根据自己的实际需要删除或者修改标签的某些类别

以KITTI数据集为例，由于实际需求，将原始的类别1和2删除，相应的类别3、4、5、6、7、8应该为类别1、2、3、4、5、6

yolo格式的KITTI数据集：

链接：https://pan.baidu.com/s/1BvXIipqmLGGe1wVS7_KWVg?pwd=q0vz
提取码：q0vz

KITTI数据集类别：

classes = {'0':'Cyclist', '1':'DontCare', '2':'Misc', '3':'Person_sitting', '4':'Tram', '5':'Truck', '6':'Van', '7':'car', '8':'person'}

修改后的类别：

classes = {'0':'Cyclist', '1':'Person_sitting', '2':'Tram', '3':'Truck', '4':'Van', '5':'car', '6':'person'}

修改代码如下：

import os
import random
import numpy as np
from numpy import *

txt_file_path = 'E:/KITTI/train/labels'  # 原始的标签路径
save_file_path = 'E:/KITTI/train/labels1'  # 修改后的标签路径

labels_name = os.listdir(txt_file_path)  # 获得每一个标签名字的列表 / os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表
num = len(labels_name)  # 获取列表长度
list = range(num)  # 创建从0到num的整数列表  list = range(0, num)
files = os.listdir(save_file_path)

for i in list:  # 遍历每一个文件
    name = labels_name[i]  # 获取每一个文件的文件名
    read_file = open(txt_file_path + "/" + name, 'r')  # 读取txt_file_path/labels路径中的文件，r表示以只读方式打开文件
    fline = read_file.readlines()  # 读取txt文件中每一行 / readlines()表示读取整行 / fline是列表类型，fline列表里的元素是str类型
    save_txt = open(save_file_path + "/" + name, 'w+')  # 读取save_file_path/labels路径中的文件. w+表示打开一个文件用于读写。如果该文件已存在则打开文件，并从开头开始编辑，即原有内容会被删除。如果该文件不存在，创建新文件。

    for j in fline:  # 遍历txt文件中每一行
        list1 = j.split()
        # print(list1)

        # 删除类别
        if list1[0] != '1':  # 删除类别1
            if list1[0] != '2':  # 删除类别2
                list2 = list1
                #print(list2)

                # 修改类别
                if list2[0] == '3':
                    list2[0] = '1'  # 将类别3改成类别1
                elif list2[0] == '4':
                    list2[0] = '2'  # 将类别4改成类别2
                elif list2[0] == '5':
                    list2[0] = '3'  # 将类别5改成类别3
                elif list2[0] == '6':
                    list2[0] = '4'  # 将类别6改成类别4
                elif list2[0] == '7':
                    list2[0] = '5'  # 将类别7改成类别5
                elif list2[0] == '8':
                    list2[0] = '6'  # 将类别8改成类别6

                b = " ".join(list2)   # 将列表转换成字符串类型，且用空格分割
                save_txt.write(b)  # 写入新的文件中
                save_txt.write('\n')  # 换行

其他数据集也一样，删除类别、更改类别、合并类别，修改 if 语句即可

批量删除并更改yolo格式的标签类别

猜你喜欢