行人跟踪数据集转为yolov5格式的数据集Labelme以及可视化

Posted 2022-11-24 AI浩

摘要

最近，在做行人检测的项目，想找一些行人的检测数据，来增强模型的效果，发现行人跟踪数据集可用，只需要把跟踪数据转为yolov5格式即可。

数据集

这些数据集来自百度飞桨的PaddleDetection项目。飞桨对一些特殊格式的数据做了格式的统一。格式如下：

Caltech
   |——————images
   |        └——————00001.jpg
   |        |—————— ...
   |        └——————0000N.jpg
   └——————labels_with_ids
            └——————00001.txt
            |—————— ...
            └——————0000N.txt
MOT17
   |——————images
   |        └——————train
   |        └——————test
   └——————labels_with_ids
            └——————train

行人跟踪

数据集	下载链接	备注
MOT17	download	-
MOT16	download	-
Caltech	download	-
Cityscapes	download	-
CUHKSYSU	download	-
PRW	download	-
ETHZ	download	-

车辆跟踪

数据集	下载链接	备注
AICity21	download	-

人头跟踪

数据集	下载链接	备注
HT21	download	-

多类别跟踪

数据集	下载链接	备注
VisDrone-MOT	download	-

数据转换

将数据集转为yolov5用的数据集

以CUHKSYSU为例，首先将其转为yolov5格式，代码如下：

import os
import shutil
import numpy as np
import configparser
import glob
if not os.path.exists('images'):
    os.makedirs('images/train')
if not os.path.exists('labels'):
    os.makedirs('labels/train')


txt_list=glob.glob('./CUHKSYSU/labels_with_ids/*.txt')
print(txt_list)
for  txt in txt_list:
    label_list=[]
    with open(txt,'r') as fs:
        txt_lines=fs.readlines()
        if len(txt_lines) <1:
            continue
        for txt_line in txt_lines:
            txt_line=txt_line.replace('\\n','').split()
            line_new=str(txt_line[0]+' '+txt_line[2]+' '+txt_line[3]+' '+txt_line[4]+' '+txt_line[5])
            label_list.append(line_new)
    txt_write=txt.replace('\\\\','/').split('/')[-1]
    image_path='./CUHKSYSU/images/'+txt_write.replace('txt','jpg')
    txt_pull_path="./labels/train/"+txt_write
    image_path_new='./images/train/'+txt_write.replace('txt','jpg')
    # image_new_path="./labels/train"+
    with open(txt_pull_path, 'w') as fw:
        for label in label_list:
            fw.write(label+'\\n')
    shutil.copy(image_path,image_path_new)

完成之后就可以看到images和labels的文件夹

可视化YoloV5数据

转换完成后，查看数据是否正确，对其做可视化操作。代码如下：

import cv2
import os

def draw_box_in_single_image(image_path, txt_path):
    # 读取图像
    image = cv2.imread(image_path)

    # 读取txt文件信息
    def read_list(txt_path):
        pos = []
        with open(txt_path, 'r') as file_to_read:
            while True:
                lines = file_to_read.readline()  # 整行读取数据
                if not lines:
                    break
                # 将整行数据分割处理，如果分割符是空格，括号里就不用传入参数，如果是逗号， 则传入‘，'字符。
                p_tmp = [float(i) for i in lines.split(' ')]
                pos.append(p_tmp)  # 添加新读取的数据
                # Efield.append(E_tmp)
                pass
        return pos


    # txt转换为box
    def convert(size, box):
        xmin = (box[1]-box[3]/2.)*size[1]
        xmax = (box[1]+box[3]/2.)*size[1]
        ymin = (box[2]-box[4]/2.)*size[0]
        ymax = (box[2]+box[4]/2.)*size[0]
        box = (int(xmin), int(ymin), int(xmax), int(ymax))
        return box

    pos = read_list(txt_path)
    print(pos)
    tl = int((image.shape[0]+image.shape[1])/2)
    lf = max(tl-1,1)
    for i in range(len(pos)):
        label = str(int(pos[i][0]))
        print('label is '+label)
        box = convert(image.shape, pos[i])
        image = cv2.rectangle(image,(box[0], box[1]),(box[2],box[3]),(0,0,255),2)
        cv2.putText(image,label,(box[0],box[1]-2), 0, 1, [0,0,255], thickness=2, lineType=cv2.LINE_AA)
        pass

    if pos:
        cv2.imwrite('./VOCData/see_images/.png'.format(image_path.split('\\\\')[-1][:-4]), image)
    else:
        print('None')

    print('./VOCData/see_images/.png'.format(image_path.split('\\\\')[-1][:-4]))
    # cv2.imshow("images", image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()


img_folder = "./images/train"
img_list = os.listdir(img_folder)
img_list.sort()

label_folder = "./labels/train"
label_list = os.listdir(label_folder)
label_list.sort()
if not os.path.exists('./VOCData/see_images'):
    os.makedirs('./VOCData/see_images')
for i in range(len(img_list)):
    image_path = img_folder + "\\\\" + img_list[i]
    txt_path = label_folder + "\\\\" + label_list[i]
    draw_box_in_single_image(image_path, txt_path)

转为Labelme格式的数据

将yolov5格式的数据转为Labelme数据格式的数据集，可以使用Labelme对其做微调和更改。代码如下：

import os
import json
import cv2
import base64

version = '3.16.7'
flags = 
lineColor = [0, 255, 0, 128]
fillColor = [255, 0, 0, 128]
filelist=os.listdir('./images/train/')
for file in filelist:
    if 'jpg' in file:
        dic = 
        dic['version'] = version
        dic['flags'] = flags
        dic['shapes'] = []
        img = cv2.imread('./images/train/'.format(file))
        imageHeight, imageWidth, _ = img.shape
        with open('./labels/train/.txt'.format(file.split('.')[0])) as f:
            datas = f.readlines()
            for data in datas:
                shape = 
                shape['label'] = 'person'
                shape['line_color'] = None
                shape['fill_color'] = None
                data = data.strip().split(' ')
                x = float(data[1]) * imageWidth
                y = float(data[2]) * imageHeight
                w = float(data[3]) * imageWidth
                h = float(data[4]) * imageHeight
                x1 = x - w / 2
                y1 = y - h / 2
                x2 = x1 + w
                y2 = y1 + h
                shape['points'] = [[x1, y1], [x2, y2]]
                shape['shape_type'] = 'rectangle'
                shape['flags'] = 
                dic['shapes'].append(shape)
        dic['lineColor'] = lineColor
        dic['fillColor'] = fillColor
        dic['imagePath'] = file
        dic['imageData'] = base64.b64encode(
            open('./images/train/'.format(file), "rb").read()).decode('utf-8')
        dic['imageHeight'] = imageHeight
        dic['imageWidth'] = imageWidth
        fw = open('./images/train/.json'.format(file.split('.')[0]), 'w')
        json.dump(dic, fw)
        fw.close()

结果如下：