目标检测脚本之yolopose 多分类标签生成

目标检测脚本之yolopose 多分类标签生成

一、需求分析

yolopose 在训练多分类模型时，需要将其它不包含关键点类别的数据标签文件整理成与人型

二、需求实现

1. 使用(0，0，0)填充例如 猫、狗、自行车等类别的17个关键点
2. 合并人型类别和之前填充的其他类别的数据

# # 一、填充非关键点标签
import os
import shutil

label_path = ""  # 未包含人型关键点的标签

label_save_path = ""  #存取填充后的标签

if not os.path.exists(label_save_path):
    os.makedirs(label_save_path)


label_files = os.listdir(label1_path)


for txt in label_files:
    txt_path = os.path.join(label1_path,txt1)

    save_txt_path = os.path.join(label_save_path,txt1)


    save_text = []

    with open(txt_path,'r') as file,open(save_txt_path,'w+') as save_file:

        #获取类别的标签记录，并用0补齐关键点
        labels = file.readlines()
        for label in labels:
            class_name = label.split(" ")[0]
            if class_name != '0':
                temp_text = ''
                for i in range(0,51):
                    temp_text += ' 0.000000'
                label_2v2 =  label.strip("\n")
                label_2v2 += temp_text
                label_2v2 += '\n'
                save_text.append(label_2v2)

        # 将填充后的标签写入
        for save_label in save_text:
            # print(save_label)
            save_file.write(save_label)

        file.close()
        save_file.close()



# 二、 整合关键点标签和非关键点标签
import os
import shutil

label1_path = ""  # 包含人型关键点的标签
label2_path = ""  # 未包含关键点的标签


label_save_path = ""  #存取整合后的标签

if not os.path.exists(label_save_path):
    os.makedirs(label_save_path)


label1_files = os.listdir(label1_path)
label2_files = os.listdir(label2_path)

count = 0
count_label1 = len(label1_files)
count_label2 = len(label2_files)

for txt1 in label1_files:
    txt1_path = os.path.join(label1_path,txt1)
    txt2_path = os.path.join(label2_path,txt1)
    
    save_txt_path = os.path.join(label_save_path,txt1)
    
    if not os.path.exists(txt2_path):  # 核对两个文件夹中是否有对应文件，label2_dir中不存在时，删除label中的
        count += 1
        print(txt1_path)
        # os.remove(txt1_path)
        # exit()
        continue
    
    
    # 包含对应的文件时，分别读取两个文件中的内容，整合后保存到第三个文件夹中
    save_text = []
    
    with open(txt1_path,'r') as file1,open(txt2_path,'r') as file2,open(save_txt_path,'w+') as save_file:
        # 获取带关键点类别的标签记录
        labels_1 = file1.readlines()
        for label_1 in labels_1:
            class_name1 = label_1.split(" ")[0]
            if class_name1 == '0':
                save_text.append(label_1)
        
        #获取其它类别的标签记录，并用0补齐关键点
        labels_2 = file2.readlines()
        for label_2 in labels_2:
            class_name2 = label_2.split(" ")[0]
            if class_name2 != '0':
                temp_text = ''
                for i in range(0,51):
                    temp_text += ' 0.000000'
                label_2v2 =  label_2.strip("\n")
                label_2v2 += temp_text
                label_2v2 += '\n'
                save_text.append(label_2v2)
        
        # 将整合的标签写入
        for save_label in save_text:
            # print(save_label)
            save_file.write(save_label)

        file1.close()
        file2.close()
        save_file.close()
    print(count_label1,count_label2,count)