【目标检测】Visdrone数据集和CARPK数据集预处理

之前的博文【目标检测】YOLOv5跑通VisDrone数据集对Visdrone数据集简介过，这里不作复述，本文主要对Visdrone数据集和CARPK数据集进行目标提取和过滤。

需求描述

本文需要将Visdrone数据集中有关车和人的数据集进行提取和合并，车标记为类别0，人标记为类别1，并转换成YOLO支持的txt格式。

Visdrone数据集

Visdrone数据集转换成YOLO的txt格式

首先对原始数据集做一个格式转换，下面这段代码延用官方提供的转换脚本。

from utils.general import download, os, Pathdef visdrone2yolo(dir):from PIL import Imagefrom tqdm import tqdmdef convert_box(size, box):# Convert VisDrone box to YOLO xywh boxdw = 1. / size[0]dh = 1. / size[1]return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh(dir / 'labels').mkdir(parents=True, exist_ok=True)  # make labels directorypbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')for f in pbar:img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).sizelines = []with open(f, 'r') as file:  # read annotation.txtfor row in [x.split(',') for x in file.read().strip().splitlines()]:if row[4] == '0':  # VisDrone 'ignored regions' class 0continuecls = int(row[5]) - 1  # 类别号-1box = convert_box(img_size, tuple(map(int, row[:4])))lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n")with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'labels' + os.sep), 'w') as fl:fl.writelines(lines)  # write label.txtdir = Path(r'E:\Dataset\VisDrone')  # datasets文件夹下Visdrone2019文件夹目录
# Convert
for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':visdrone2yolo(dir / d)  # convert VisDrone annotations to YOLO labels

标签可视化

对txt标签进行可视化，查看过滤之前的效果。

import os
import numpy as np
import cv2# 修改输入图片文件夹
img_folder = "image"
img_list = os.listdir(img_folder)
img_list.sort()
# 修改输入标签文件夹
label_folder = "labels2"
label_list = os.listdir(label_folder)
label_list.sort()
# 输出图片文件夹位置
path = os.getcwd()
output_folder = path + '/' + str("output")
os.mkdir(output_folder)# 坐标转换
def xywh2xyxy(x, w1, h1, img):label, x, y, w, h = x# print("原图宽高:\nw1={}\nh1={}".format(w1, h1))# 边界框反归一化x_t = x * w1y_t = y * h1w_t = w * w1h_t = h * h1# print("反归一化后输出：\n第一个:{}\t第二个:{}\t第三个:{}\t第四个:{}\t\n\n".format(x_t, y_t, w_t, h_t))# 计算坐标top_left_x = x_t - w_t / 2top_left_y = y_t - h_t / 2bottom_right_x = x_t + w_t / 2bottom_right_y = y_t + h_t / 2# print('标签:{}'.format(labels[int(label)]))# print("左上x坐标:{}".format(top_left_x))# print("左上y坐标:{}".format(top_left_y))# print("右下x坐标:{}".format(bottom_right_x))# print("右下y坐标:{}".format(bottom_right_y))# 绘制矩形框# cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), colormap[1], 2)# (可选)给不同目标绘制不同的颜色框if int(label) == 0:cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), (0, 255, 0), 2)elif int(label) == 1:cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), (255, 0, 0), 2)else:cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), (0, 0, 0), 2)return imgif __name__ == '__main__':for i in range(len(img_list)):image_path = img_folder + "/" + img_list[i]label_path = label_folder + "/" + label_list[i]# 读取图像文件img = cv2.imread(str(image_path))h, w = img.shape[:2]# 读取 labelswith open(label_path, 'r') as f:lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)# 绘制每一个目标for x in lb:# 反归一化并得到左上和右下坐标，画出矩形框img = xywh2xyxy(x, w, h, img)"""# 直接查看生成结果图cv2.imshow('show', img)cv2.waitKey(0)"""cv2.imwrite(output_folder + '/' + '{}.png'.format(image_path.split('/')[-1][:-4]), img)

可视化效果如图所示：
注：该数据集对人的姿态还进行区分，行走状态的人划分为pedestrian，其它姿态(比如躺下或坐下)标记为people。

在这里插入图片描述

过滤标签

具体过滤规则：

合并car、van、truck、bus为car(0)
合并pedestrian，people为person(1)
舍弃其它类别

import os
import numpy as np
from tqdm import tqdm# Visdrone类别
# names: ['pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor' ]# 修改输入标签文件夹
label_folder = "labels"
label_list = os.listdir(label_folder)# 标签输出文件夹
label_output = "labels2"# class_set
car_set = [3, 4, 5, 8]
person_set = [0, 1]if __name__ == '__main__':for label_file in tqdm(os.listdir(label_folder)):# 读取 labelswith open(os.path.join(label_folder, label_file), 'r') as f:lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)# 写入 labelswith open(os.path.join(label_output, label_file), 'a') as f:for obj in lb:# 若是行人，修改类别为1if int(obj[0]) in person_set:obj[0] = 1f.write(('%g ' * 5).rstrip() % tuple(obj) + '\n')# 若是车辆，修改类别为0elif int(obj[0]) in car_set:obj[0] = 0f.write(('%g ' * 5).rstrip() % tuple(obj) + '\n')

过滤之后的效果如图所示：

在这里插入图片描述

CARPK数据集

CARPK数据集是无人机在40米高空拍摄的汽车数据集，里面仅包含汽车单一目标。

下载地址：https://github.com/zstar1003/Dataset

原始label格式：

1019 521 1129 571 1
1013 583 1120 634 1

对应含义为： xmin, ymin, xmax, ymax，cls

处理脚本：

import os
import numpy as np
from tqdm import tqdm# 修改输入标签文件夹
# label_folder = r"E:\Dataset\CARPK_devkit\data\Annotations"
label_folder = r"annotations"
label_list = os.listdir(label_folder)# 标签输出文件夹
label_output = r"labels"# 图像宽高
img_width = 1280
img_height = 720if __name__ == '__main__':for label_file in tqdm(os.listdir(label_folder)):# 读取 labelswith open(os.path.join(label_folder, label_file), 'r') as f:lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=int)for obj in lb:class_index = obj[4]xmin, ymin, xmax, ymax = obj[0], obj[1], obj[2], obj[3]# 将box信息转换到yolo格式xcenter = xmin + (xmax - xmin) / 2ycenter = ymin + (ymax - ymin) / 2w = xmax - xminh = ymax - ymin# 绝对坐标转相对坐标，保存6位小数xcenter = round(xcenter / img_width, 6)ycenter = round(ycenter / img_height, 6)w = round(w / img_width, 6)h = round(h / img_height, 6)info = [str(i) for i in [class_index, xcenter, ycenter, w, h]]# 写入 labelswith open(os.path.join(label_output, label_file), 'a') as f:# 若文件不为空，添加换行if os.path.getsize(os.path.join(label_output, label_file)):f.write("\n" + " ".join(info))else:f.write(" ".join(info))