|
- #import torch.utils.data as data
-
- from PIL import Image
- import os
- import os.path
- import numpy as np
- from numpy.random import randint
-
- import mindspore.dataset as ds
-
- from distributed_sampler import DistributedSampler
-
- 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 UCF101DataSet():
- def __init__(self, root_path, list_file, mode = 'train',
- num_segments=3, new_length=1, modality='RGB',
- image_tmpl='img_{:05d}.jpg', transform=None,
- force_grayscale=False, random_shift=True, test_mode=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
-
- if self.modality == 'RGBDiff':
- self.new_length += 1# Diff needs one more image to calculate diff
-
- self._parse_list()
-
- def _load_image(self, directory, idx):
- if self.modality == 'RGB' or self.modality == 'RGBDiff':
- return [Image.open(os.path.join(directory, self.image_tmpl.format(idx))).convert('RGB')]
- elif self.modality == 'Flow':
- x_img = Image.open(os.path.join(directory, self.image_tmpl.format('x', idx))).convert('L')
- y_img = Image.open(os.path.join(directory, self.image_tmpl.format('y', idx))).convert('L')
-
- return [x_img, y_img]
-
- def _parse_list(self):
- self.video_list = [VideoRecord(x.strip().split(' ')) for x in open(self.list_file)]
-
- def _sample_indices(self, record):
- """
-
- :param record: VideoRecord
- :return: list
- """
-
- 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_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)])
-
- return offsets + 1
-
- def __getitem__(self, index):
- record = self.video_list[index]
-
- if not self.test_mode:
- segment_indices = self._sample_indices(record) if self.random_shift else self._get_val_indices(record)
- else:
- segment_indices = self._get_test_indices(record)
-
- return self.get(record, segment_indices)
-
- 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)
-
- def Create_UCF101_Dataset(data_path, datalist_path, mode='train', batch_size = 16, shuffle = True, device_num = 1, is_distributed = False, rank = 0):
-
- dataset = UCF101DataSet(data_path, datalist_path)
- current_sampler = None
- if(is_distributed):
- distributed_sampler = DistributedSampler(len(dataset), device_num, rank, shuffle=shuffle)
- current_sampler = distributed_sampler
- else:
- current_sampler = DistributedSampler(len(dataset), device_num, rank, shuffle=shuffle)
- #current_sampler = SequentialSampler(start_index=0, num_samples=len(dataset))
- #hwc_to_chw = CV.HWC2CHW()
- op_none = lambda x: x
-
- data_set = ds.GeneratorDataset(dataset, column_names=["frames","label"], shuffle=True,sampler=current_sampler)
- data_set = data_set.map(input_columns=["frames"], operations=op_none, num_parallel_workers=8)
- data_set = data_set.map(input_columns=["label"], operations=op_none, num_parallel_workers=8)
- data_set = data_set.batch(batch_size, drop_remainder=True)
-
- return data_set
|