txts os.listdir(label_dir) # 所有的txt文件 如本人数据集 {list: 1080} [ 20210819B000001.txt , ...]
data_list [] # 存放所有txt文件的name 去掉.txt txt中各个类别的数量 data_list append(temp)
for txt in txts: # 遍历每一个txt文件(yolo格式)
cnt_list [0] * class_num # 初始化[0, 0, 0] 存放所有txt文件中每一个类别的数量
temp [] # 存放当前txt文件的name 去掉.txt 当前txt中各个类别的数量 如(3类) [ 20210819B000001 , 0, 1, 0]
name txt.split( . )[0] # 获得当前txt文件的name 如 20210819B000001
temp.append(name) # 当前txt文件的name加入temp中
with open(os.path.join(label_dir, txt), r ) as f: # 打开当前txt文件 读取每一行信息 存放在lines中
lines f.readlines()
for line in lines: # 遍历当前txt文件中的每一行信息 line class x y w h
label int(line.split()[0]) # 得到当前行的class信息
if label_start_num 1:
cnt_list[label - 1] 1
elif label_start_num 0:
cnt_list[label] 1
temp.extend(cnt_list)
data_list.append(temp)
# 过滤出只含指定类别的文件名
data np.array(data_list)
index np.where(data[:, select_num 1], True, False)
return data[index][:, [0, select_num 1]]
def getSum(data):
sum 0
for d in data:
sum int(d[1])
return sum
def plot_class_distribution(data):
x [ 0:白粉病 , 1:潜叶蛾 , 2:炭疽病 ]
fig plt.figure(figsize (10, 8)) # 画布大小和像素密度
plt.bar(x, data, width 0.5, align center )
for a, b, i in zip(x, data, range(len(x))): # zip 函数
plt.text(a, b 0.01, %d % int(data[i]), ha center , fontsize 15, color r ) # plt.text 函数
plt.xticks(fontsize 15)
plt.yticks(fontsize 15)
plt.xlabel( 类别 , fontsize 16)
plt.ylabel( 数量 , fontsize 16)
plt.title( 广佛手病虫害训练集各类别分布情况 , fontsize 16)
plt.show()
# 保存到本地
# plt.savefig( )
