【ReID】【代码注释】读数据/吐数据 deep-person-reid/dataset_loader.py

1,286 阅读1分钟

源码URL: github.com/michuanhaoh…

读数据/吐数据,读前39行源码并注释

from __future__ import print_function, absolute_import
import os
from PIL import Image
import numpy as np
import os.path as osp

import torch  # 需要重构pytorch的dataloader函数
from torch.utils.data import Dataset



def read_image(img_path):  # 定义读图片函数,传img_path进去
    got_img = False  # 定义标志位,判断是否读到图片以便报错

    # 如果图片有问题?如果网络通信断了?如果因为各种原因没读到图片?
    if not os.path.exists(img_path):  # path是否存在
        raise IOError("{} dose not exist".format(img_path))  # 不存在则报错"img_path dose not exist"
    while not got_img:  # 若没有读到图片,while循环
        try:  # 尝试读图片
            img = Image.open(img_path).convert('RGB')  # 函数核心,打开图片,转RGB
            got_img = True  # 读到图片,标志位变为True
        except IOError:  # 否则出现异常
            print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))  # 报错
            pass  # 循环跳出
    return img  # 返回img


class ImageDataset(Dataset):  # 旨在重构pytorch的dataset,将Dataset作为参数放入,继承Dataset的类
    def __init__(self, dataset, transform=None):  # 传入dataset,目前不考虑数据增广
        self.dataset = dataset  # 定义dataset
        self.transform = transform  # 要用到,先定义为类的属性

    def __len__(self):  # 在直接返回dataset的长度,返回图片数量
        return len(self.dataset)

    def __getitem__(self, index):
        img_path, pid, camid = self.dataset[index]  # 结合data_manager的dataset,得到三个参数
        img = read_image(img_path)  # 用刚刚的function把img读出来
        if self.transform is not None:  # 如果有transform
            img = self.transform  # 执行
        return img, pid, camid  # 直接返回img和pid,camid




if __name__ == '__main__':
    import data_manager
    dataset = data_manager.init_img_dataset(root='F:/Market-1501/Market-1501-v15.09.15', name='market1501')
    train_loader = ImageDataset(dataset.train)
    from IPython import embed
    embed()
    """
    In [1]: for batch_id, (img, pid, camid) in enumerate(train_lodaer):
       ...: 	break
       ...: 
    In [2]: img
    Out[2]: <PIL.Image.Image image mode=RGB size=64x128 at 0x16C128F5A58>
    In [3]: img.save('aaaa.jpg')
    """

保存图像aaaa.jpg结果查看:

aaaa.jpg