VisDrone2019数据集提取行人并转为voc数据集

 参考了一下别人的代码，自己琢磨着写了下。

import os
import csv
import xmltodict 
from PIL import Image

def get_box(box):
    xmin = box[0] 
    ymin = box[1]
    xmax = box[0] + box[2]
    ymax = box[1] + box[3]
    return (xmin,ymin,xmax,ymax)

if __name__ == "__main__":
    person_number = 0
    txt_path = r'E:\MBY\small\VisDrone2019-DET-test\test1'
    im_path = r'E:\MBY\small\VisDrone2019-DET-test\images'
    save_path = r'E:\MBY\small\VisDrone2019-DET-test\test'

    for i in os.listdir(txt_path):
        object_ = []
        with open(os.path.join(txt_path,i), newline='') as csvfile:
            spamreader = csv.reader(csvfile)
            for row in spamreader:
                if (int(row[5]) - 1) < 2:
                    person_number+=1
                    if row[4] == '0':
                        continue
                    bb = get_box(tuple(map(int, row[:4])))
                    mess = {'name': 'person',
                            'pose': 'Unspecified',
                            'truncated': '0',
                            'difficult': '0',
                            'bndbox': {'xmin':str(bb[0]), 'ymin':str(bb[1]), 'xmax':str(bb[2]), 'ymax':str(bb[3])}}
                    object_.append(mess)
        if object_ != []:
            with Image.open(os.path.join(im_path,i.replace('txt','jpg'))) as Img:
                img_size = Img.size
                Img.save(os.path.join(save_path,i.replace('txt','jpg')))
            mess_all = {
                'annotation': {
                    'folder': 'JPEGImages',
                    'filename': i.replace('txt','jpg'),
                    'path': os.path.join(im_path,i.replace('txt','jpg')),
                    'source': {
                        'database': 'Unknown'
                    },
                    'size': {
                        'width': str(img_size[0]),
                        'height': str(img_size[1]),
                        'depth': '3'
                    },
                    'segmented': '0',
                    'object': object_
                }
            }
            mess_all
            xml_str = xmltodict.unparse(mess_all, encoding='utf-8')
            f = open(os.path.join(save_path,i.replace('txt','xml')), 'w', encoding="UTF-8")
            f.write(xml_str)
            f.close()    

 

参考：

• https://blog.csdn.net/weixin_41170972/article/details/117044387