保姆级教程：用Python脚本将VisDrone2019数据集一键转成COCO格式（附完整代码）

从VisDrone到COCO无人机视觉数据格式转换实战指南当你第一次打开VisDrone2019数据集时可能会被那些密密麻麻的.txt标注文件弄得一头雾水。作为计算机视觉领域最常用的无人机视角数据集之一VisDrone却采用了与主流框架不兼容的自定义格式——这就像拿到了珍贵的食材却找不到合适的菜谱。本文将带你用Python脚本打通这个关键环节把VisDrone2019-DET数据集完美转换为COCO格式让你能立即在MMDetection、Detectron2等框架中开展目标检测实验。1. 理解VisDrone与COCO格式的本质差异在动手写代码之前我们需要先搞清楚两种格式的DNA差异。VisDrone2019-DET数据集采用简单的每图对应一个.txt文件的标注方式而COCO格式则将所有标注信息整合到一个结构化的JSON文件中。VisDrone标注文件示例例如000001.txt641,420,73,72,0,4,1,1 322,424,97,51,0,4,1,1 ...每行代表一个目标实例各字段含义为前4个数字边界框的左上角x,y坐标和宽高x,y,w,h第5个数字目标可见性评分0-1第6个数字类别ID对应11类物体最后2个数字暂时不用关注相比之下COCO格式的核心结构包含三个关键部分{ images: [{id: 1, file_name: 000001.jpg, ...}], annotations: [{image_id: 1, category_id: 4, bbox: [x,y,w,h], ...}], categories: [{id: 1, name: pedestrian}, ...] }2. 项目环境搭建与目录准备开始转换前确保你的Python环境已安装以下关键包pip install opencv-python tqdm numpyVisDrone2019-DET的标准目录结构应该是这样的VisDrone2019-DET/ ├── VisDrone2019-DET-train/ │ ├── annotations/ # 训练集标注文件 │ └── images/ # 训练集图片 ├── VisDrone2019-DET-val/ │ ├── annotations/ # 验证集标注文件 │ └── images/ # 验证集图片 └── VisDrone2019-DET-test-dev/ # 测试集无标注提示建议在输出目录预先创建好annotations/子目录这是MMDetection等框架默认寻找COCO标注文件的位置。3. 核心转换代码逐行解析让我们构建一个完整的转换脚本visdrone2coco.py关键函数如下3.1 类别映射定义categories [ {id: 0, name: ignored regions}, {id: 1, name: pedestrian}, {id: 2, name: people}, {id: 3, name: bicycle}, # ...完整类别见原始代码 ]VisDrone的11个类别需要与COCO格式中的categories数组精确对应。注意第一个类别是ignored regions这在评估时需要特殊处理。3.2 主转换函数def convert_to_coco(dir_path, output_dir): # 初始化COCO数据结构 coco_dict { images: [], annotations: [], categories: categories } annotation_id 0 # 每个标注实例的唯一ID for mode in [train, val]: img_dir os.path.join(dir_path, fVisDrone2019-DET-{mode}, images) ann_dir os.path.join(dir_path, fVisDrone2019-DET-{mode}, annotations) for filename in tqdm(os.listdir(ann_dir)): # 处理每个标注文件 base_name filename.replace(.txt, ) img_path os.path.join(img_dir, base_name .jpg) # 获取图片尺寸 img cv2.imread(img_path) height, width img.shape[:2] # 添加图片信息 coco_dict[images].append({ id: base_name, file_name: base_name .jpg, height: height, width: width }) # 解析标注文件 with open(os.path.join(ann_dir, filename)) as f: for line in f: parts [int(x) for x in line.strip().split(,)] if len(parts) 6: continue # 构建COCO标注格式 bbox [parts[0], parts[1], parts[2], parts[3]] coco_dict[annotations].append({ id: annotation_id, image_id: base_name, category_id: parts[5], bbox: bbox, area: bbox[2] * bbox[3], iscrowd: 0, segmentation: [] }) annotation_id 1 # 保存为JSON文件 os.makedirs(output_dir, exist_okTrue) for mode in [train, val]: subset_dict { images: [img for img in coco_dict[images] if img[file_name].startswith(mode[:3])], annotations: [ann for ann in coco_dict[annotations] if ann[image_id].startswith(mode[:3])], categories: categories } with open(f{output_dir}/instances_{mode}2019.json, w) as f: json.dump(subset_dict, f)3.3 边界框处理技巧VisDrone的边界框有时会超出图像范围我们需要进行裁剪处理def clamp_bbox(bbox, img_width, img_height): x, y, w, h bbox x1 max(0, x) y1 max(0, y) x2 min(img_width, x w) y2 min(img_height, y h) return [x1, y1, x2 - x1, y2 - y1]4. 实战应用与验证转换完成后我们可以用以下代码验证生成的COCO格式文件是否正确from pycocotools.coco import COCO import matplotlib.pyplot as plt # 加载生成的COCO标注 coco COCO(output/instances_train2019.json) # 随机可视化一个样本 img_id coco.getImgIds()[0] img_info coco.loadImgs(img_id)[0] ann_ids coco.getAnnIds(imgIdsimg_id) anns coco.loadAnns(ann_ids) img cv2.imread(os.path.join(VisDrone2019-DET-train/images, img_info[file_name])) plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) coco.showAnns(anns) plt.show()5. 进阶技巧与性能优化对于大规模数据集转换可以考虑以下优化策略多进程处理使用Python的multiprocessingfrom multiprocessing import Pool def process_file(args): filename, ann_dir, img_dir args # ...处理单个文件的逻辑 with Pool(processes4) as pool: args_list [(f, ann_dir, img_dir) for f in os.listdir(ann_dir)] results pool.map(process_file, args_list)内存优化分批处理图片避免同时加载所有图像数据使用生成器逐步构建JSON结构格式扩展# 添加额外的COCO字段可选 annotation.update({ attributes: { occluded: parts[4] 1, # 根据可见性评分判断是否遮挡 weather: sunny # 可从文件名推断天气条件 } })6. 常见问题解决方案问题1转换后某些标注框位置异常检查确认是否正确处理了边界框越界情况修复添加clamp_bbox函数确保坐标在图像范围内问题2COCO验证时报类别ID不连续原因VisDrone的类别ID从0开始但0对应的是忽略区域解决在categories数组中确保ID连续性问题3大型数据集转换内存不足方案改用流式JSON写入方式import ijson def stream_write_json(output_path, data_generator): with open(output_path, w) as f: f.write({images: [) # 分批写入图片数据 f.write(], annotations: [) # 分批写入标注数据 f.write(], categories: [...]})7. 集成到训练流程转换后的数据集可以直接用于主流框架。以MMDetection为例配置文件中只需指定data dict( traindict( typeCocoDataset, ann_filedata/VisDrone/annotations/instances_train2019.json, img_prefixdata/VisDrone/VisDrone2019-DET-train/images/, classesclasses # 你的类别列表 ), valdict(...), testdict(...) )对于需要YOLO格式的情况可以基于COCO格式进一步转换# 将COCO的[x,y,w,h]转换为YOLO的[center_x, center_y, w, h]归一化 def coco2yolo(bbox, img_w, img_h): x_center (bbox[0] bbox[2]/2) / img_w y_center (bbox[1] bbox[3]/2) / img_h width bbox[2] / img_w height bbox[3] / img_h return [x_center, y_center, width, height]