【WiderPerson数据集处理成yolov5数据集格式】

【WiderPerson数据集处理成yolov5数据集格式】

1.下载数据集

2.处理数据集

1）只保留有annotation的数据

2）只保留上述有1和3标签的

3）把坐标转化成yolov5格式

# -*- ecoding: utf-8 -*-
# @ModuleName: 
# @Author: 
# @Time: 
import os
import cv2
import shutil
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

if __name__ == "__main__":
    img_dir = 'F:/e/data/human_detect/widerperson/Images/'
    anno_dir = 'F:/e/data/human_detect/widerperson/Annotations/'

    img_save_dir = 'F:/e/data/human_detect/widerperson/select/Images/'
    anno_save_dir = 'F:/e/data/human_detect/widerperson/select/Annotations/'

    if not os.path.exists(img_save_dir):
        os.makedirs(img_save_dir)
    if not os.path.exists(anno_save_dir):
        os.makedirs(anno_save_dir)

    img_list = os.listdir(img_dir)
    img_list.sort()

    label_num = [1, 3]
    for i, idx in enumerate(img_list):
        img_path = os.path.join(img_dir, idx)
        img_name = idx.split('.')[0]
        anno_path = os.path.join(anno_dir, img_name + '.jpg.txt')
        img = cv2.imread(img_path)
        h = img.shape[0]
        w = img.shape[1]

        if os.path.isfile(anno_path):
            anno_file = open(anno_path, 'r')
            anno_lines = anno_file.readlines()
            num = 0
            anno_save_path = os.path.join(anno_save_dir, img_name + '.txt')
            save_anno = open(anno_save_path, 'w')
            # widerperson官方txt第一行是标签数量
            for j, jdx in enumerate(anno_lines[1:]):
                label = int(jdx.split(' ')[0])
                if label in label_num:
                    if label == 1:
                        label = 0
                    elif label == 3:
                        label = 1
                    xmin = float(jdx.split(' ')[1])
                    ymin = float(jdx.split(' ')[2])
                    xmax = float(jdx.split(' ')[3])
                    ymax = float(jdx.split(' ')[4].split('\n')[0])

                    box = (xmin, xmax, ymin, ymax)
                    # 转换成yolov5格式
                    bb = convert((w, h), box)
                    save_anno.write(str(label) + " " + " ".join([str(a) for a in bb]) + '\n')

                    num += 1

            save_anno.close()
            img_save_path = os.path.join(img_save_dir, idx)
            shutil.copy(img_path, img_save_path)
            if num == 0:
                os.remove(img_save_path)
                os.remove(anno_save_path)

3.划分数据集

# -*- ecoding: utf-8 -*-
# @ModuleName: 
# @Author: 
# @Time: 
import os
import shutil
import random

if __name__ == "__main__":
    random.seed(0)

    img_dir = 'F:/e/data/human_detect/widerperson/select/Images/'
    anno_dir = 'F:/e/data/human_detect/widerperson/select/Annotations/'

    train_img_dir ='F:/e/data/human_detect/widerperson/dataset/images/train/'
    train_label_dir = 'F:/e/data/human_detect/widerperson/dataset/labels/train/'
    val_img_dir = 'F:/e/data/human_detect/widerperson/dataset/images/val/'
    val_label_dir = 'F:/e/data/human_detect/widerperson/dataset/labels/val/'
    test_img_dir = 'F:/e/data/human_detect/widerperson/dataset/images/test/'
    test_label_dir = 'F:/e/data/human_detect/widerperson/dataset/labels/test/'
    if not os.path.exists(train_img_dir):
        os.makedirs(train_img_dir)
    if not os.path.exists(train_label_dir):
        os.makedirs(train_label_dir)
    if not os.path.exists(val_img_dir):
        os.makedirs(val_img_dir)
    if not os.path.exists(val_label_dir):
        os.makedirs(val_label_dir)
    if not os.path.exists(test_img_dir):
        os.makedirs(test_img_dir)
    if not os.path.exists(test_label_dir):
        os.makedirs(test_label_dir)

    trainval_percent = 1
    train_percent = 0.9

    img_list = os.listdir(img_dir)
    img_list.sort()

    num = len(img_list)
    list = range(num)

    tv = int(num * trainval_percent)
    tr = int(tv * train_percent)
    trainval = random.sample(list, tv)
    train = random.sample(trainval, tr)

    print("train and val size", tv)
    print("train size", tr)

    for i in list:
        name = img_list[i].split('.')[0]

        if i in trainval:
            if i in train:
                shutil.copy(os.path.join(img_dir, name + '.jpg'), os.path.join(train_img_dir, name + '.jpg'))
                shutil.copy(os.path.join(anno_dir, name + '.txt'), os.path.join(train_label_dir, name + '.txt'))
            else:
                shutil.copy(os.path.join(img_dir, name + '.jpg'), os.path.join(val_img_dir, name + '.jpg'))
                shutil.copy(os.path.join(anno_dir, name + '.txt'), os.path.join(val_label_dir, name + '.txt'))
        else:
            shutil.copy(os.path.join(img_dir, name + '.jpg'), os.path.join(test_img_dir, name + '.jpg'))
            shutil.copy(os.path.join(anno_dir, name + '.txt'), os.path.join(test_label_dir, name + '.txt'))

参考博客

1.【数据集处理】WiderPerson介绍以及转YOLO格式（图片教程及代码----超详细）

2.睿智的目标检测56——Pytorch搭建YoloV5目标检测平台