Pytorch 创建自己的数据DataSet 并加载多张图像数据到input tensor

#DataLoader.py

import torch
from torch.utils.data import Dataset,DataLoader
import numpy as np
import os
import torch.utils.data as data
from PIL import Image
import cv2
import os.path
from numpy.random import randint
import torchvision
import glob
import random
import time
from transforms import *
#存储数据列表文档所在目录
'''
 video_datasets -  'category.txt'
                -  'train_videofolder.txt'
                -  'val_videofolder.txt'
'''
ROOT_DATASET = 'video_datasets'
def return_ucf101(modality):
    filename_categories = 'category.txt'
    if modality == 'RGB':
        #图像所在目录
        root_data = '/home/enbo/share/UCF101'
        #root_data = '/mnt/localssd1/bzhou/something/20bn-something-something-v1'
        filename_imglist_train = 'train_videofolder.txt'
        filename_imglist_val = 'val_videofolder.txt'
        prefix = '{:04d}.jpg'
    elif modality == 'Flow':
        root_data = '/data/vision/oliva/scratch/bzhou/video/something-something/flow'
        #root_data = '/mnt/localssd1/bzhou/something/flow'
        filename_imglist_train = 'something/train_videofolder.txt'
        filename_imglist_val = 'something/val_videofolder.txt'
        prefix = '{:05d}.jpg'
    else:
        print('no such modality:' + modality)
        os.exit()
    return filename_categories, filename_imglist_train, filename_imglist_val, root_data, prefix

def return_dataset_list(dataset,modality):
    dataset_dict = {'UCF101':return_ucf101}
    if dataset in dataset_dict:
        file_categories,file_imglist_train,file_imglist_val,root_data,prefix = dataset_dict[dataset](modality)
    else:
        raise ValueError('Unknown dataset ' + dataset)
    file_imglist_train = os.path.join(ROOT_DATASET, file_imglist_train)
    file_imglist_val = os.path.join(ROOT_DATASET, file_imglist_val)
    file_categories = os.path.join(ROOT_DATASET, file_categories)
    with open(file_categories) as f:
        lines = f.readlines()
    categories = [item.rstrip() for item in lines]
    return categories, file_imglist_train, file_imglist_val, root_data, prefix

class VideoRecord(object):
    def __init__(self, row):
        self._data = row

    @property
    def path(self):
        return self._data[0]

    @property
    def num_frames(self):
        return int(self._data[1])

    @property
    def label(self):
        return int(self._data[2])


class MydataSet(data.Dataset):

    def __init__(self, root_path, list_file, num_segment =3,
                 new_length =1, modality='RGB', image_tmpl='img_{:05d}.jpg', transform=None,
                 random_shift=True, test_mode=False, remove_missing=False, dense_sample=False):
        self.root_path = root_path
        self.list_file = list_file
        self.num_segments = num_segments
        self.new_length = new_length
        self.modality = modality
        self.image_tmpl = image_tmpl
        self.transform = transform
        self.random_shift = random_shift
        self.test_mode = test_mode
        self.remove_missing = remove_missing
        self.dense_sample = dense_sample  # using dense sample as I3D

        self._parse_list()
        self.total_list = self.video_list

    def _load_image(self, directory, idx):
        if self.modality == 'RGB' or self.modality == 'RGBDiff':
           # print(self.root_path)
           # print(self.image_tmpl.format(idx))
            try:
                return [
                    cv2.imread(os.path.join(self.root_path, directory, directory + '-' + self.image_tmpl.format(idx)))[:, :,
                    ::-1]]
                # return [Image.open(os.path.join(self.root_path, directory, directory + '-' + self.image_tmpl.format(idx))).convert('RGB')]
            except Exception:
                print('error loading image:',
                      os.path.join(self.root_path, directory, directory + '-' + self.image_tmpl.format(idx)))
                raise
                # return [Image.open(os.path.join(self.root_path, directory, self.image_tmpl.format(1))).convert('RGB')]
        elif self.modality == 'Flow':
            try:
                idx_skip = 1 + (idx - 1) * 5
                flow = Image.open(os.path.join(self.root_path, directory, self.image_tmpl.format(idx_skip))).convert('RGB')
            except Exception:
                print('error loading flow file:', os.path.join(self.root_path, directory, self.image_tmpl.format(idx_skip)))
                flow = Image.open(os.path.join(self.root_path, directory, self.image_tmpl.format(1))).convert('RGB')
            # the input flow file is RGB image with (flow_x, flow_y, blank) for each channel
            flow_x, flow_y, _ = flow.split()
            x_img = flow_x.convert('L')
            y_img = flow_y.convert('L')

            return [x_img, y_img]

    def _parse_list(self):
        #过滤掉不足八张的子目录
        tmp =[x.strip().split(' ') for x in open(self.list_file)]
        tmp = [item for item in tmp if int(item[1]) >= 8]
        self.video_list = [VideoRecord(item) for item in tmp]
      #  print(self.video_list[0])  #把每一个图像目转换成对象
        print('video number:%d'%(len(self.video_list)))
    #dataset 内置的方法，通过该方法，读取

    def _get_val_indices(self, record):
        if record.num_frames > self.num_segments + self.new_length - 1:
            tick = (record.num_frames - self.new_length + 1) / float(self.num_segments)
            offsets = np.array([int(tick / 2.0 + tick * x) for x in range(self.num_segments)])
        else:
            offsets = np.zeros((self.num_segments,))

        return offsets + 1

    def _get_test_indices(self, record):
        tick = (record.num_frames - self.new_length + 1) / float(self.num_segments)

        offsets = np.array([int(tick / 2.0 + tick * x) for x in range(self.num_segments)])

    def __getitem__(self,index):

        record = self.video_list[index]
        while not os.path.exists(os.path.join(self.root_path, record.path, record.path + '-' + self.image_tmpl.format(1))):
           # print('[DEBUG:getitem]')
            print(os.path.join(self.root_path, record.path, record.path + '-' + self.image_tmpl.format(1)))
            index = np.random.randint(len(self.video_list))
            record = self.video_list[index]
        if not self.test_mode:
           # print('[DEBUG:getitem]')
           # print(record.path)
            segment_indices = self._sample_indices(record) if self.random_shift else self._get_val_indices(record)
            print(segment_indices)
        else:
            segment_indices = self._get_test_indices(record)

        return self.get(record, segment_indices)

    def _sample_indices(self, record):
        #print(record.path)
        #print(record.num_frames)
        average_duration = (record.num_frames - self.new_length + 1) // self.num_segments
        if average_duration > 0:
            offsets = np.multiply(list(range(self.num_segments)), average_duration) + randint(average_duration,
                                                                                              size=self.num_segments)
        elif record.num_frames > self.num_segments:
            offsets = np.sort(randint(record.num_frames - self.new_length + 1, size=self.num_segments))
        else:
            offsets = np.zeros((self.num_segments,))
        return offsets + 1

    def get(self, record, indices):
        images = list()
        for seg_ind in indices:
            p = int(seg_ind)
            for i in range(self.new_length):

                seg_imgs = self._load_image(record.path, p)
                #向列表追加值
                images.extend(seg_imgs)
                if p < record.num_frames:
                    p += 1

        process_data = self.transform(images)
        return process_data, record.label

    def __len__(self):
        return len(self.video_list)


#类别      训练目录               图像所在路径
categories, train_list, val_list, root_path, prefix = return_dataset_list('UCF101', 'RGB')

data_length = 1
num_segments = 8
input_mean = [0.485, 0.456, 0.406]
input_std = [0.229, 0.224, 0.225]
normalize = GroupNormalize(input_mean, input_std)
arch = 'RESBET'
#测试
# a = MydataSet(root_path, train_list, num_segments,
#           new_length=data_length,
#           modality="RGB",
#           image_tmpl=prefix,
#           transform=torchvision.transforms.Compose([
#               GroupMultiScaleCrop(224, [1, .875, .75, .66]),
#               GroupRandomHorizontalFlip(is_flow=False),
#               Stack(roll=(arch in ['BNInception', 'InceptionV3'])),
#               ToTorchFormatTensor(div=(arch not in ['BNInception', 'InceptionV3'])),
#               normalize,
#           ]))
# img,label = a[0]
# print(img.size)
#transform 用法；tran=transform.compose([]) tran(img)
batch_size = 1
workers = 2
train_loader = torch.utils.data.DataLoader(
               MydataSet(root_path, train_list, num_segments,
               new_length = data_length,
               modality="RGB",
               image_tmpl=prefix,
               transform=torchvision.transforms.Compose([
                  GroupMultiScaleCrop(224, [1, .875, .75, .66]),
                  GroupRandomHorizontalFlip(is_flow=False),
                  Stack(roll=(arch in ['BNInception', 'InceptionV3'])),
                  ToTorchFormatTensor(div=(arch not in ['BNInception', 'InceptionV3'])),
                  normalize,
              ])),
    batch_size=batch_size, shuffle=True,
    num_workers=workers, pin_memory=True,
    drop_last=True)  # prevent something not % n_GPU

for i, (input, target) in enumerate(train_loader):
    print(input.size())  #
    break

'''
  总结：
      DataSet 数据加载步骤
       一： 通过初始化把train_folder.txt中的内容中的每一行变成一个object,然后所有的存放在一个list中 每个obj包含子文件夹路径，图像个数，label
       
       二：通过 __getItem__这个函数每调用一次就返回num_segment个图像的数据
       通过transform 转换成tensor  [8,224,224,3](可以变)

'''

import torchvision
import random
from PIL import Image, ImageOps
import numpy as np
import numbers
import math
import torch
import cv2


class GroupRandomCrop(object):
    def __init__(self, size):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size

    def __call__(self, img_group):

        w, h = img_group[0].size
        th, tw = self.size

        out_images = list()

        x1 = random.randint(0, w - tw)
        y1 = random.randint(0, h - th)

        for img in img_group:
            assert(img.size[0] == w and img.size[1] == h)
            if w == tw and h == th:
                out_images.append(img)
            else:
                out_images.append(img.crop((x1, y1, x1 + tw, y1 + th)))

        return out_images


class GroupCenterCrop(object):
    def __init__(self, size):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
        # self.worker = torchvision.transforms.CenterCrop(size)

    def __call__(self, img_group):
        # return [self.worker(img) for img in img_group]
        return [cv2.resize(img, self.size) for img in img_group]

#随机旋转
class GroupRandomHorizontalFlip(object):
    """Randomly horizontally flips the given PIL.Image with a probability of 0.5
    """
    def __init__(self, is_flow=False):
        self.is_flow = is_flow

    def __call__(self, img_group, is_flow=False):
        v = random.random()
        if v < 0.5:
            # ret = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_group]
            ret = [cv2.flip(img, 1) for img in img_group]
            if self.is_flow:
                for i in range(0, len(ret), 2):
                    ret[i] = ImageOps.invert(ret[i])  # invert flow pixel values when flipping
            return ret
        else:
            return img_group


class GroupNormalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
        rep_mean = self.mean * (tensor.size()[0]//len(self.mean))
        rep_std = self.std * (tensor.size()[0]//len(self.std))

        # TODO: make efficient
        for t, m, s in zip(tensor, rep_mean, rep_std):
            t.sub_(m).div_(s)

        return tensor


class GroupScale(object):
    """ Rescales the input PIL.Image to the given 'size'.
    'size' will be the size of the smaller edge.
    For example, if height > width, then image will be
    rescaled to (size * height / width, size)
    size: size of the smaller edge
    interpolation: Default: PIL.Image.BILINEAR
    """

    def __init__(self, size, interpolation=Image.BILINEAR):
        self.worker = torchvision.transforms.Resize(size, interpolation)

    def __call__(self, img_group):
        return [self.worker(img) for img in img_group]


class GroupOverSample(object):
    def __init__(self, crop_size, scale_size=None):
        self.crop_size = crop_size if not isinstance(crop_size, int) else (crop_size, crop_size)

        if scale_size is not None:
            self.scale_worker = GroupScale(scale_size)
        else:
            self.scale_worker = None

    def __call__(self, img_group):

        if self.scale_worker is not None:
            img_group = self.scale_worker(img_group)

        image_w, image_h = img_group[0].size
        crop_w, crop_h = self.crop_size

        offsets = GroupMultiScaleCrop.fill_fix_offset(False, image_w, image_h, crop_w, crop_h)
        oversample_group = list()
        for o_w, o_h in offsets:
            normal_group = list()
            flip_group = list()
            for i, img in enumerate(img_group):
                crop = img.crop((o_w, o_h, o_w + crop_w, o_h + crop_h))
                normal_group.append(crop)
                flip_crop = crop.copy().transpose(Image.FLIP_LEFT_RIGHT)

                if img.mode == 'L' and i % 2 == 0:
                    flip_group.append(ImageOps.invert(flip_crop))
                else:
                    flip_group.append(flip_crop)

            oversample_group.extend(normal_group)
            oversample_group.extend(flip_group)
        return oversample_group

#是否裁剪
class GroupMultiScaleCrop(object):

    def __init__(self, input_size, scales=None, max_distort=1, fix_crop=True, more_fix_crop=True):
        self.scales = scales if scales is not None else [1, .875, .75, .66]
        self.max_distort = max_distort
        self.fix_crop = fix_crop
        self.more_fix_crop = more_fix_crop
        self.input_size = input_size if not isinstance(input_size, int) else [input_size, input_size]
        self.interpolation = cv2.INTER_LINEAR

    def __call__(self, img_group):

        #im_size = img_group[0].size
        im_size = img_group[0].shape[:2]


        crop_w, crop_h, offset_w, offset_h = self._sample_crop_size(im_size)
       # print(crop_w, crop_h, offset_w, offset_h)
        # crop_img_group = [img.crop((offset_w, offset_h, offset_w + crop_w, offset_h + crop_h)) for img in img_group]
        crop_img_group = [img[offset_w: offset_w + crop_w, offset_h: offset_h + crop_h] for img in img_group]
        #ret_img_group = [img.resize((self.input_size[0], self.input_size[1]), self.interpolation) for img in crop_img_group]
        ret_img_group = [cv2.resize(img, (self.input_size[0], self.input_size[1]), interpolation=self.interpolation) for img in crop_img_group]
        return ret_img_group

    def _sample_crop_size(self, im_size):
        image_w, image_h = im_size[0], im_size[1]

        # find a crop size
        base_size = min(image_w, image_h)
        crop_sizes = [int(base_size * x) for x in self.scales]
        crop_h = [self.input_size[1] if abs(x - self.input_size[1]) < 3 else x for x in crop_sizes]
        crop_w = [self.input_size[0] if abs(x - self.input_size[0]) < 3 else x for x in crop_sizes]

        pairs = []
        for i, h in enumerate(crop_h):
            for j, w in enumerate(crop_w):
                if abs(i - j) <= self.max_distort:
                    pairs.append((w, h))

        crop_pair = random.choice(pairs)
        if not self.fix_crop:
            w_offset = random.randint(0, image_w - crop_pair[0])
            h_offset = random.randint(0, image_h - crop_pair[1])
        else:
            w_offset, h_offset = self._sample_fix_offset(image_w, image_h, crop_pair[0], crop_pair[1])

        return crop_pair[0], crop_pair[1], w_offset, h_offset

    def _sample_fix_offset(self, image_w, image_h, crop_w, crop_h):
        offsets = self.fill_fix_offset(self.more_fix_crop, image_w, image_h, crop_w, crop_h)
        return random.choice(offsets)

    @staticmethod
    def fill_fix_offset(more_fix_crop, image_w, image_h, crop_w, crop_h):
        w_step = (image_w - crop_w) // 4
        h_step = (image_h - crop_h) // 4

        ret = list()
        ret.append((0, 0))  # upper left
        ret.append((4 * w_step, 0))  # upper right
        ret.append((0, 4 * h_step))  # lower left
        ret.append((4 * w_step, 4 * h_step))  # lower right
        ret.append((2 * w_step, 2 * h_step))  # center

        if more_fix_crop:
            ret.append((0, 2 * h_step))  # center left
            ret.append((4 * w_step, 2 * h_step))  # center right
            ret.append((2 * w_step, 4 * h_step))  # lower center
            ret.append((2 * w_step, 0 * h_step))  # upper center

            ret.append((1 * w_step, 1 * h_step))  # upper left quarter
            ret.append((3 * w_step, 1 * h_step))  # upper right quarter
            ret.append((1 * w_step, 3 * h_step))  # lower left quarter
            ret.append((3 * w_step, 3 * h_step))  # lower righ quarter

        return ret


class GroupRandomSizedCrop(object):
    """Random crop the given PIL.Image to a random size of (0.08 to 1.0) of the original size
    and and a random aspect ratio of 3/4 to 4/3 of the original aspect ratio
    This is popularly used to train the Inception networks
    size: size of the smaller edge
    interpolation: Default: PIL.Image.BILINEAR
    """
    def __init__(self, size, interpolation=Image.BILINEAR):
        self.size = size
        self.interpolation = interpolation

    def __call__(self, img_group):
        for attempt in range(10):
            area = img_group[0].size[0] * img_group[0].size[1]
            target_area = random.uniform(0.08, 1.0) * area
            aspect_ratio = random.uniform(3. / 4, 4. / 3)

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

            if random.random() < 0.5:
                w, h = h, w

            if w <= img_group[0].size[0] and h <= img_group[0].size[1]:
                x1 = random.randint(0, img_group[0].size[0] - w)
                y1 = random.randint(0, img_group[0].size[1] - h)
                found = True
                break
        else:
            found = False
            x1 = 0
            y1 = 0

        if found:
            out_group = list()
            for img in img_group:
                img = img.crop((x1, y1, x1 + w, y1 + h))
                assert(img.size == (w, h))
                out_group.append(img.resize((self.size, self.size), self.interpolation))
            return out_group
        else:
            # Fallback
            scale = GroupScale(self.size, interpolation=self.interpolation)
            crop = GroupRandomCrop(self.size)
            return crop(scale(img_group))


class Stack(object):

    def __init__(self, roll=False):
        self.roll = roll

    def __call__(self, img_group):
        if self.roll:
            #print(self.roll)
            return np.concatenate([np.array(x)[:, :, ::-1] for x in img_group], axis=2)
        else:
            # return np.concatenate(img_group, axis=2)
            # return num_segments * H * W * C
            return np.array(img_group)


class ToTorchFormatTensor(object):
    """ Converts a PIL.Image (RGB) or numpy.ndarray (H x W x C) in the range [0, 255]
    to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] """
    def __init__(self, div=True):
        self.div = div

    def __call__(self, pic):
        if isinstance(pic, np.ndarray):
            print(pic.shape)
            # handle numpy array
            #img = torch.from_numpy(pic).permute(2, 0, 1).contiguous()
            img = torch.from_numpy(pic).permute(3, 0, 1, 2).contiguous()
        else:
            # handle PIL Image
            img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
            img = img.view(pic.size[1], pic.size[0], len(pic.mode))
            # put it from HWC to CHW format
            # yikes, this transpose takes 80% of the loading time/CPU
            img = img.transpose(0, 1).transpose(0, 2).contiguous()
        return img.float().div(255) if self.div else img.float()


class IdentityTransform(object):

    def __call__(self, data):
        return data


# if __name__ == "__main__":
#     trans = torchvision.transforms.Compose([
#         GroupScale(256),
#         GroupRandomCrop(224),
#         Stack(),
#         ToTorchFormatTensor(),
#         GroupNormalize(
#             mean=[.485, .456, .406],
#             std=[.229, .224, .225]
#         )]
#     )
#
#     im = Image.open('../tensorflow-model-zoo.torch/lena_299.png')
#
#     color_group = [im] * 3
#     rst = trans(color_group)
#
#     gray_group = [im.convert('L')] * 9
#     gray_rst = trans(gray_group)
#
#     trans2 = torchvision.transforms.Compose([
#         GroupRandomSizedCrop(256),
#         Stack(),
#         ToTorchFormatTensor(),
#         GroupNormalize(
#             mean=[.485, .456, .406],
#             std=[.229, .224, .225])
#     ])
#     print(trans2(color_group))