Save the model detection results as gt to the annotation file

Due to the needs of the final design, the label information of the "person" category should be added to the previously marked vehicle pictures, but considering the time cost, consider using a ready-made training model to identify the people in the picture as gt.

1.yolov5

• yolov5: https://github.com/ultralytics/yolov5
• Use the yolov5 algorithm to detect the "person" appearing in the count picture, and record the bbox
• Run detect.py directly, with its own category filtering, if –classes is set to 0, only the person category will be saved.

for *xyxy, conf, cls in reversed(det):
    if save_txt:  # Write to file
        bbox = torch.tensor(xyxy).tolist()
        with open(f'{txt_path}.txt', 'a') as f:
             f.write(str(bbox[0]) + " " + str(bbox[1]) + " " + str(bbox[2]) + " " + str(bbox[3]) + '\n')

2. Modify the corresponding xml

The original data set is in voc format, and the corresponding "person" is added to the corresponding xml
Related xml operations

1) Modify the label name

import xml.etree.ElementTree as ET
import os
ann_filepath="ann路径"
def edit_xml(xml_file):
    tree = ET.parse(xml_file)
    objs = tree.findall('object')
    for ix, obj in enumerate(objs):
        name = obj.find('name')
        if name.text == 'Taxi':
            name.text = "Car"
        elif name.text == 'Other Car':
            name.text = "Car"
    tree.write(xml_file, method='xml', encoding='utf-8')  # 更新xml文件

if __name__ == '__main__':
    for xml in os.listdir(ann_filepath):
        edit_xml(ann_filepath + xml)  # 修改xml文件

2) Add object tag

import xml.etree.ElementTree as ET
import os

ann_filepath="xml路径"
txt_flepath="txt路径"
def edit_xml(xml_file,txt_file):
    tree = ET.parse(xml_file)
    rootdir = tree.getroot()
    if os.path.exists(txt_file):
        f1 = open(txt_file,'r',encoding= 'UTF-8')
        for line in f1:
            data = line.split(" ")
            object = ET.Element("object")
            name = ET.Element("name")
            name.text="Person"
            occluded = ET.Element("occluded")
            occluded.text = "0"
            bndbox = ET.Element("bndbox")
            xmin = ET.Element("xmin")
            xmin.text=str(data[0])
            ymin = ET.Element("ymin")
            ymin.text=str(data[1])
            xmax = ET.Element("xmax")
            xmax.text=str(data[2])
            ymax = ET.Element("ymax")
            ymax.text=str(data[3])
            bndbox.append(xmin)
            bndbox.append(ymin)
            bndbox.append(xmax)
            bndbox.append(ymax)
            object.append(name)
            object.append(occluded)
            object.append(bndbox)
            rootdir.append(object)
        tree.write(xml_file, method='xml', encoding='utf-8')  
if __name__ == '__main__':
    for xml in os.listdir(ann_filepath):
        name = xml.split(".")[0]+'.txt'
        edit_xml(ann_filepath + xml,txt_flepath+name)  # 修改xml文件

3) Generate xml based on txt

import os
ann_filepath="ann保存路径"
txt_filepath="txt路径"
from lxml import etree
def creat_xml(xml_file,txt_file,imgname):
    root = etree.Element("annotation")  # 创建根节点 annotation
    folder = etree.SubElement(root, "folder")
    folder.text = "frame"
    filename = etree.SubElement(root, "filename")
    filename.text = imgname

    source = etree.SubElement(root, "source")
    database = etree.SubElement(source, "database")
    database.text="Unknown"
    annotation1 = etree.SubElement(source, "annotation")
    annotation1.text="annotation"
    image = etree.SubElement(source, "image")
    image.text="Unknown"

    size = etree.SubElement(root, "size")
    width = etree.SubElement(size, "width")
    width.text="1920"
    height = etree.SubElement(size, "height")
    height.text="1080"
    depth = etree.SubElement(size, "depth")

    if os.path.exists(txt_file):
        f1 = open(txt_file, 'r', encoding='UTF-8')
        for line in f1:
            data = line.split(" ")
            object = etree.SubElement(root, "object")  # 创建annotation节点的子节点 object
            namen = etree.SubElement(object, "name")  # 创建object节点的子节点 name
            namen.text = 'Person'  # name的值为person
            occludedn = etree.SubElement(object, "occluded")
            occludedn.text="0"
            bndbox = etree.SubElement(object, "bndbox")
            xminn = etree.SubElement(bndbox, "xmin")
            xminn.text = str(data[0])
            yminn = etree.SubElement(bndbox, "ymin")
            yminn.text = str(data[1])
            xmaxn = etree.SubElement(bndbox, "xmax")
            xmaxn.text = str(data[2])
            ymaxn = etree.SubElement(bndbox, "ymax")
            ymaxn.text = str(data[3])
    tree = etree.ElementTree(root)
    tree.write(xml_file, pretty_print=True, xml_declaration=False, encoding='utf-8')


if __name__ == '__main__':
    for txt in os.listdir(txt_filepath):
        name = txt.split(".")[0]+'.xml'
        imgname = txt.split(".")[0]+'.jpg'
        creat_xml(ann_filepath + name,txt_filepath+txt,imgname)

3. Generate train\test.txt according to the xml file

import os
import random
import argparse
parser = argparse.ArgumentParser()
ann_filepath="Annotations路径"
txt_file="ImageSets/Main/ txt保存路径"
# 1.全部文件的文件夹路径，根据自己的数据进行修改。
parser.add_argument('--totalfiles_path', default=ann_filepath, type=str, help='input xml label path')
# 2.保存 划分后生成的txt文件 的文件夹路径。
parser.add_argument('--txt_path', default=txt_file, type=str, help='output txt label path')
opt = parser.parse_args()

needfiles_percent = 0.8 # 3.占总文件的百分比
totalfilespath = opt.totalfiles_path
txtsavepath = opt.txt_path
total_path = os.listdir(totalfilespath)

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

num = len(total_path)
list_index = range(num)
nf = int(num * needfiles_percent)
needfiles = random.sample(list_index, nf)


file_need = open(txtsavepath + '/person_train.txt', 'w')
file_remain = open(txtsavepath + '/person_val.txt', 'w')


for i in list_index:
    name = total_path[i][:-4] + '\n' # 4.-4可以控制txt文件中文件名要不要加上后缀名
    if i in needfiles:
        file_need.write(name)
    else:
        file_remain.write(name)

file_need.close()
file_remain.close()