!转载请注明原文地址!——东方旅行者
更多行人重识别文章移步我的专栏:行人重识别专栏
数据加载器(dataset_loader.py)
一、数据加载器作用
dataset_loader类将dataset_manager生成的元素为(图片地址,行人ID,摄像机ID)的索引列表数据集转换为元素为(图片,行人ID,摄像机ID)的图片数据集。
(图片地址,行人ID,摄像机ID)——>(图片,行人ID,摄像机ID)
二、数据加载器编写思路
- 引用dataset_manager,dataset_manager产生的三个数据集(训练集、测试集、查询集)
- ImageDataset类根据索引列表产生图片数据集,ImageDataset类继承torch.utils下data中的Dataset,继承Dataset类必须重写其中的__len__与__getitem__方法,前者获取数据集的大小,后者方法根据索引获取数据集中对应的数据。在循环遍历ImageDataset类实例化的对象时,Torch会根据这个getitem方法遍历对象中的每个数据。
- ImageDataset类中的getitem方法调用一个read_image方法,该方法根据图片地址读取图片,该图片的类型为PIL.Image.Image,最终dataset_loader产生的数据集中每条数据为一个三元组(图片,行人ID,摄像机ID)。
三、代码
import os
from PIL import Image
import numpy as np
import os.path as osp
import torch
from torch.utils.data import Dataset
"""
ImageDataset类用于
将dataset_manager生成的元素为(图片地址,行人ID,摄像机ID)的索引列表数据集转换为元素为(图片,行人ID,摄像机ID)的图片数据集
即:(图片地址,行人ID,摄像机ID)——>(图片,行人ID,摄像机ID)
"""
def read_image(img_path):
if not osp.exists(img_path):
raise IOError("{} does not exist.".format(img_path))
else:
img=Image.open(img_path).convert('RGB')
return img
class ImageDataset(Dataset):
"""
所有继承torch.utils.data.Dataset类都需要重写其中的__len__与__getitem__方法
前者用于获取数据集的大小
后者用于数据集迭代过程中获取其中每一条数据
"""
def __init__(self,dataset):
self.dataset=dataset
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
img_path, pid, cid=self.dataset[index]
img=read_image(img_path)
return img, pid, cid
if __name__=='__main__':
from dataset_manager import Market1501
dataset=Market1501()
train_loader=ImageDataset(dataset.train)
from IPython import embed
for batch_id, (img, pid, cid) in enumerate(train_loader):
print(batch_id,img,pid,cid)