Pascal VOC 数据集导入并转为yolo格式

将voc数据集转换为yolo格式，yolo格式是[x，y,长，宽，类别]

首先通过使用一个名为 VOCDetection 的类来加载 Pascal VOC 数据集的一个子集，我导入的是2007版本的，‘train’数据集也可以是‘val’（验证集）、‘test’（测试集）等。如果 download=True 并且数据集已经存在，torchvision 不会重新下载，而是会直接加载现有数据集。如果 download=False 并且数据集已经存在，torchvision 同样会加载现有数据集。

如果确定数据集已经下载并且位于正确的位置，可以安全地将 download 参数设置为 False 以避免不必要的网络检查。这通常是在多次运行代码或者在一个不允许网络访问的环境中运行代码时的做法。

from torchvision.datasets import VOCDetection

dataset = VOCDetection(root = './data', year = '2007', image_set = 'train', download = False)

成功显示

Dataset VOCDetection
    Number of datapoints: 2501
    Root location: ./data

查看dataset[0]

(<PIL.Image.Image image mode=RGB size=500x333>,
 {'annotation': {'folder': 'VOC2007',
   'filename': '000012.jpg',
   'source': {'database': 'The VOC2007 Database',
    'annotation': 'PASCAL VOC2007',
    'image': 'flickr',
    'flickrid': '207539885'},
   'owner': {'flickrid': 'KevBow', 'name': '?'},
   'size': {'width': '500', 'height': '333', 'depth': '3'},
   'segmented': '0',
   'object': [{'name': 'car',
     'pose': 'Rear',
     'truncated': '0',
     'difficult': '0',
     'bndbox': {'xmin': '156', 'ymin': '97', 'xmax': '351', 'ymax': '270'}}]}})

创建一个字典，方便之后的类别判断。初始将所有类别设置为0。

# 创建一个字典，其中键是类别名称，值是对应的类别ID
VOC_CLASSES = [
    'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
    'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
    'dog', 'horse', 'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor'
]

voc_class_dict = {class_name: 0 for class_name in VOC_CLASSES}

首先创建一个yolo空列表，之后逐渐添加数据。提取dataset数据，之后遍历所有数据，如果name在VOC_CLASSES中则相应的name的数据由零变为一。（注意每次循环之前都要重新将字典值转换成0，不然逐渐遍历之后，类别数值都会转换为一），之后提取obj每个数据的x，y最值，x，y值分布为x，y的最大值-最小值的二分之一。宽w的值为x最大值-最小值。高h的值为x的最大值-最小值。之后创建字典d_l存储x,y,w,h值，并将VOC_CLASSES的值转换为列表格式，将两个列表合并，之后yolo列表添加d_l，return yolo_data

def yolo(dataset):
    # 创建一个空列表来存储所有YOLO格式的数据
    yolo_data = []
    for da in metadata:
        # 提取字典
        annotation_dict = da

        # 提取object列表
        objects = annotation_dict['annotation']['object']

       # 遍历每个object，打印其信息
        for obj in objects:
            VOC_CLASSES = [
                 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
                 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
                 'dog', 'horse', 'motorbike', 'person', 'pottedplant',
                 'sheep', 'sofa', 'train', 'tvmonitor'
            ]
            voc_class_dict = {class_name: 0 for class_name in VOC_CLASSES}
            name = obj['name']
            bndbox = obj['bndbox']
            xmin = int(bndbox['xmin'])
            ymin = int(bndbox['ymin'])
            xmax = int(bndbox['xmax'])
            ymax = int(bndbox['ymax'])
            x = (xmin + xmax)/2
            y = (ymin + ymax)/2
            w = xmax - xmin
            h = ymax - ymin
            if name in voc_class_dict:
                voc_class_dict[name] = 1
            d_l = [x,y,w,h]
            class_values = list(voc_class_dict.values())
            d_l = d_l  + class_values
            # 将YOLO格式的数据添加到列表中
            yolo_data.append(d_l)
    return yolo_data

准备数据，筛选voc数据中每个样本的数据

metadata = [value[1] for value in dataset]

结果[{'annotation': {'folder': 'VOC2007',
   'filename': '000012.jpg',
   'source': {'database': 'The VOC2007 Database',
    'annotation': 'PASCAL VOC2007',
    'image': 'flickr',
    'flickrid': '207539885'},
   'owner': {'flickrid': 'KevBow', 'name': '?'},
   'size': {'width': '500', 'height': '333', 'depth': '3'},
   'segmented': '0',
   'object': [{'name': 'car',
     'pose': 'Rear',
     'truncated': '0',
     'difficult': '0',
     'bndbox': {'xmin': '156', 'ymin': '97', 'xmax': '351', 'ymax': '270'}}]}}, ......]

调用函数

yolo_data = yolo(metadata)
print(yolo_data)

结果[[253.5, 183.5, 195, 173, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [232.0, 130.5, 94, 137, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], .......]