【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目标检测平台