行人重识别(9)——代码实践之数据加载器（dataset_loader.py）

！转载请注明原文地址！——东方旅行者

更多行人重识别文章移步我的专栏：行人重识别专栏

数据加载器（dataset_loader.py）

一、数据加载器作用

dataset_loader类将dataset_manager生成的元素为（图片地址，行人ID，摄像机ID）的索引列表数据集转换为元素为（图片，行人ID，摄像机ID）的图片数据集。
（图片地址，行人ID，摄像机ID）——>（图片，行人ID，摄像机ID）

二、数据加载器编写思路

1. 引用dataset_manager，dataset_manager产生的三个数据集（训练集、测试集、查询集）
2. ImageDataset类根据索引列表产生图片数据集，ImageDataset类继承torch.utils下data中的Dataset，继承Dataset类必须重写其中的__len__与__getitem__方法，前者获取数据集的大小，后者方法根据索引获取数据集中对应的数据。在循环遍历ImageDataset类实例化的对象时，Torch会根据这个getitem方法遍历对象中的每个数据。
3. 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)

四、测试结果