在用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 语句即可