pytorch读取VOC数据集
简单介绍VOC数据集
首先介绍下VOC2007数据集(下图是VOC数据集格式，为了叙述方便，我这里只放了两张图像)
Main文件夹内的trainval.txt中的内容如下：存储了图像的名称不加后缀。
000009
000052
Annotations中存储的是标注文件，以xml文件存储。这里简单截个图说明一下：
<annotation>
<folder>VOC2007</folder>
<filename>000009.jpg</filename> # ！存储图像名称，若转换voc数据集，记住命名格式最好是相同
<source>
<database>The VOC2007 Database</database>
<annotation>PASCAL VOC2007</annotation>
<image>flickr</image>
<flickrid>325443404</flickrid>
</source>
<owner>
<flickrid>autox4u</flickrid>
<name>Perry Aidelbaum</name>
</owner>
<size> # 存储着图像的宽和高以及深度
<width>500</width>
<height>375</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object> # 存储着对象
<name>horse</name>
<pose>Right</pose>
<truncated>0</truncated>
<difficult>0</difficult> # 是否是困难检测目标
<bndbox> # 图像中的box
<xmin>69</xmin>
<ymin>172</ymin>
<xmax>270</xmax>
<ymax>330</ymax>
</bndbox>
</object>
</annotation>
pytorch读取脚本
然后你可以通过以下脚本进行读取–>（读取脚本在SSD中的方式也挺好，这里仅仅提供一种最mini的读取方式）。
import torch
import xml.etree.ElementTree as ET
import os
import cv2
import numpy as np
from torchvision import transforms
class VOCDataset(torch.utils.data.Dataset):
CLASSES_NAME = (
"__background__ ", # 记得加上背景类
"pottedplant",
"person",
"horse",
)
# 初始化类
def __init__(self, root_dir, resize_size=[800, 1024], split='trainval', use_difficult=False):
self.root = root_dir
self.use_difficult = use_difficult
self.imgset = split
self._annopath = os.path.join(self.root, "Annotations", "%s.xml")
self._imgpath = os.path.join(self.root, "JPEGImages", "%s.jpg")
self._imgsetpath = os.path.join(self.root, "ImageSets", "Main", "%s.txt")
# 读取trainval.txt中内容
with open(self._imgsetpath % self.imgset) as f: # % 是python字符串中的一个转义字符可以百度下，不难
self.img_ids = f.readlines()
self.img_ids = [x.strip() for x in self.img_ids] # ['000009', '000052']
self.name2id = dict(zip(VOCDataset.CLASSES_NAME, range(len(VOCDataset.CLASSES_NAME))))
self.resize_size = resize_size
self.mean = [0.485, 0.456, 0.406] # voc数据集中所有图像矩阵的均值和方差，为后续图像归一化做准备
self.std = [0.229, 0.224, 0.225]
print("INFO=====>voc dataset init finished ! !")
def __len__(self):
return len(self.img_ids)
def _read_img_rgb(self, path):
return cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)
def __getitem__(self, index):
img_id = self.img_ids[index]
img = self._read_img_rgb(self._imgpath % img_id)
anno = ET.parse(self._annopath % img_id).getroot() # 读取xml文档的根节点
boxes = []
classes = []
for obj in anno.iter("object"):
difficult = int(obj.find("difficult").text) == 1
if not self.use_difficult and difficult:
continue
_box = obj.find("bndbox")
box = [
_box.find("xmin").text,
_box.find("ymin").text,
_box.find("xmax").text,
_box.find("ymax").text,
]
TO_REMOVE = 1 # 由于像素是网格存储，坐标2实质表示第一个像素格，所以-1
box = tuple(
map(lambda x: x - TO_REMOVE, list(map(float, box)))
)
boxes.append(box)
name = obj.find("name").text.lower().strip()
classes.append(self.name2id[name]) # 将类别映射回去
boxes = np.array(boxes, dtype=np.float32)
#将img,box和classes转成tensor
img = transforms.ToTensor()(img) # transforms 自动将 图像进行了归一化，
boxes = torch.from_numpy(boxes)
classes = torch.LongTensor(classes)
return img, boxes, classes
if __name__ == '__main__':
dataset = VOCDataset('E://Z_summary_net/Read_VOC/VOCdevkit/VOC2007/') # 实例化一个对象
img,box,cls = dataset[0] # 返回第一张图像及box和对应的类别
print(img.shape)
print(box)
print(cls)
# 这里简单做一下可视化
# 由于opencv读入是矩阵，而img现在是tensor，因此，首先将tensor转成numpy.array
img_ = (img.numpy()*255).astype(np.uint8).transpose(1,2,0)# 注意由于图像像素分布0-255，所以转成uint8
print(img_.shape)
cv2.imshow('test',img_)
cv2.waitKey(0)
注意
这里有一个细节：torch将box、classes和img转换成张量tensor。假如现在我想利用OpenCV重新还原出图像来，那么存在一个问题：就是经过transforms.ToTensor()的img已经被进行了归一化，img的像素矩阵除以了255。因此，要想还原cv2中的numpy数组，需要*255。另外，由于cv2需要读取[W,H,C]的矩阵，因此，还需要通过transpose(1,2,0)交换一下维度。
最后，贴上transforms.ToTensor()的源码：
if isinstance(pic, np.ndarray):
# handle numpy array
if pic.ndim == 2:
pic = pic[:, :, None]
img = torch.from_numpy(pic.transpose((2, 0, 1))) # 交换维度
# backward compatibility
if isinstance(img, torch.ByteTensor): # /255 归一化
return img.float().div(255)
else:
return img