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

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

需求描述

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

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

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

from utils.general import download, os, Path def visdrone2yolo(dir): from PIL import Image from tqdm import tqdm def convert_box(size, box): # Convert VisDrone box to YOLO xywh box dw = 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 directory pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}') for f in pbar: img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size lines = [] with open(f, 'r') as file: # read annotation.txt for row in [x.split(',') for x in file.read().strip().splitlines()]: if row[4] == '0': # VisDrone 'ignored regions' class 0 continue cls = int(row[5]) - 1 # 类别号-1 box = 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.txt dir = 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 * w1 y_t = y * h1 w_t = w * w1 h_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 / 2 top_left_y = y_t - h_t / 2 bottom_right_x = x_t + w_t / 2 bottom_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 img if __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] # 读取 labels with 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)): # 读取 labels with 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) # 写入 labels with open(os.path.join(label_output, label_file), 'a') as f: for obj in lb: # 若是行人，修改类别为1 if int(obj[0]) in person_set: obj[0] = 1 f.write(('%g ' * 5).rstrip() % tuple(obj) + '\n') # 若是车辆，修改类别为0 elif int(obj[0]) in car_set: obj[0] = 0 f.write(('%g ' * 5).rstrip() % tuple(obj) + '\n')

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

CARPK数据集

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

下载地址： github /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 = 720 if __name__ == '__main__': for label_file in tqdm(os.listdir(label_folder)): # 读取 labels with 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) / 2 ycenter = ymin + (ymax - ymin) / 2 w = xmax - xmin h = 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]] # 写入 labels with 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))

可视化验证转换效果：