# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import sys import cv2 import math import random import functools try: import cPickle as pickle from cStringIO import StringIO except ImportError: import pickle from io import BytesIO import numpy as np import paddle import paddle.fluid as fluid try: from nvidia.dali.pipeline import Pipeline import nvidia.dali.ops as ops import nvidia.dali.types as types import tempfile from nvidia.dali.plugin.paddle import DALIGenericIterator except: Pipeline = object print("DALI is not installed, you can improve performance if use DALI") from PIL import Image, ImageEnhance import logging from reader_utils import DataReader logger = logging.getLogger(__name__) python_ver = sys.version_info ucf101_root = "/ssd4/chaj/ucf101/" class VideoRecord(object): ''' define a class method which used to describe the frames information of videos 1. self._data[0] is the frames' path 2. self._data[1] is the number of frames 3. self._data[2] is the label of frames ''' def __init__(self, row): self._data = row @property def path(self): return ucf101_root + "/rawframes/" + self._data[0] @property def num_frames(self): return int(self._data[1]) @property def label(self): return int(self._data[2]) class UCF101Reader(DataReader): """ Data reader for kinetics dataset of two format mp4 and pkl. 1. mp4 or avi, the original format of kinetics400 2. pkl, the mp4 or avi was decoded previously and stored as pkl 3. frames, the mp4 or avi was decoded previously and stored as frames In all cases, load the data, and then get the frame data in the form of numpy and label as an integer. dataset cfg: format num_classes seg_num short_size target_size num_reader_threads buf_size image_mean image_std batch_size list """ def __init__(self, name, mode, cfg): super(UCF101Reader, self).__init__(name, mode, cfg) self.format = cfg.MODEL.format self.num_classes = self.get_config_from_sec('model', 'num_classes') self.seg_num = self.get_config_from_sec('model', 'seg_num') self.seglen = self.get_config_from_sec('model', 'seglen') self.seg_num = self.get_config_from_sec(mode, 'seg_num', self.seg_num) self.short_size = self.get_config_from_sec(mode, 'short_size') self.target_size = self.get_config_from_sec(mode, 'target_size') self.num_reader_threads = self.get_config_from_sec(mode, 'num_reader_threads') self.buf_size = self.get_config_from_sec(mode, 'buf_size') self.fix_random_seed = self.get_config_from_sec(mode, 'fix_random_seed') self.img_mean = np.array(cfg.MODEL.image_mean).reshape( [3, 1, 1]).astype(np.float32) self.img_std = np.array(cfg.MODEL.image_std).reshape( [3, 1, 1]).astype(np.float32) # set batch size and file list self.batch_size = cfg[mode.upper()]['batch_size'] self.filelist = cfg[mode.upper()]['filelist'] # set num_trainers and trainer_id when distributed training is implemented self.num_trainers = self.get_config_from_sec(mode, 'num_trainers', 1) self.trainer_id = self.get_config_from_sec(mode, 'trainer_id', 0) self.use_dali = self.get_config_from_sec(mode, 'use_dali', False) self.dali_mean = cfg.MODEL.image_mean * (self.seg_num * self.seglen) self.dali_std = cfg.MODEL.image_std * (self.seg_num * self.seglen) if self.mode == 'infer': self.video_path = cfg[mode.upper()]['video_path'] else: self.video_path = '' if self.fix_random_seed: random.seed(0) np.random.seed(0) self.num_reader_threads = 1 def create_reader(self): # if use_dali to improve performance if self.use_dali: return self.build_dali_reader() # if set video_path for inference mode, just load this single video if (self.mode == 'infer') and (self.video_path != ''): # load video from file stored at video_path _reader = self._inference_reader_creator( self.video_path, self.mode, seg_num=self.seg_num, seglen=self.seglen, short_size=self.short_size, target_size=self.target_size, img_mean=self.img_mean, img_std=self.img_std) else: assert os.path.exists(self.filelist), \ '{} not exist, please check the data list'.format( self.filelist) _reader = self._reader_creator( self.filelist, self.mode, seg_num=self.seg_num, seglen=self.seglen, short_size=self.short_size, target_size=self.target_size, img_mean=self.img_mean, img_std=self.img_std, shuffle=(self.mode == 'train'), num_threads=self.num_reader_threads, buf_size=self.buf_size, format=self.format) def _batch_reader(): batch_out = [] for imgs, label in _reader(): if imgs is None: continue batch_out.append((imgs, label)) if len(batch_out) == self.batch_size: yield batch_out batch_out = [] return _batch_reader def _inference_reader_creator(self, video_path, mode, seg_num, seglen, short_size, target_size, img_mean, img_std): def reader(): try: imgs = mp4_loader(video_path, seg_num, seglen, mode) if len(imgs) < 1: logger.error('{} frame length {} less than 1.'.format( video_path, len(imgs))) yield None, None except: logger.error('Error when loading {}'.format(mp4_path)) yield None, None imgs_ret = imgs_transform(imgs, mode, seg_num, seglen, short_size, target_size, img_mean, img_std) label_ret = video_path yield imgs_ret, label_ret return reader def _reader_creator(self, pickle_list, mode, seg_num, seglen, short_size, target_size, img_mean, img_std, shuffle=False, num_threads=1, buf_size=1024, format='avi'): def decode_mp4(sample, mode, seg_num, seglen, short_size, target_size, img_mean, img_std): sample = sample[0].split(' ') mp4_path = ucf101_root + "/videos/" + sample[0] + ".avi" # when infer, we store vid as label label = int(sample[1]) - 1 try: imgs = mp4_loader(mp4_path, seg_num, seglen, mode) if len(imgs) < 1: logger.error('{} frame length {} less than 1.'.format( mp4_path, len(imgs))) return None, None except: logger.error('Error when loading {}'.format(mp4_path)) return None, None return imgs_transform( imgs, mode, seg_num, seglen, short_size, target_size, img_mean, img_std, name=self.name), label def decode_pickle(sample, mode, seg_num, seglen, short_size, target_size, img_mean, img_std): pickle_path = sample[0] try: if python_ver < (3, 0): data_loaded = pickle.load(open(pickle_path, 'rb')) else: data_loaded = pickle.load( open(pickle_path, 'rb'), encoding='bytes') vid, label, frames = data_loaded if len(frames) < 1: logger.error('{} frame length {} less than 1.'.format( pickle_path, len(frames))) return None, None except: logger.info('Error when loading {}'.format(pickle_path)) return None, None if mode == 'train' or mode == 'valid' or mode == 'test': ret_label = label elif mode == 'infer': ret_label = vid imgs = video_loader(frames, seg_num, seglen, mode) return imgs_transform( imgs, mode, seg_num, seglen, short_size, target_size, img_mean, img_std, name=self.name), ret_label def decode_frames(sample, mode, seg_num, seglen, short_size, target_size, img_mean, img_std): recode = VideoRecord(sample[0].split(' ')) frames_dir_path = recode.path # when infer, we store vid as label label = recode.label try: imgs = frames_loader(recode, seg_num, seglen, mode) if len(imgs) < 1: logger.error('{} frame length {} less than 1.'.format( frames_dir_path, len(imgs))) return None, None except: logger.error('Error when loading {}'.format(frames_dir_path)) return None, None return imgs_transform( imgs, mode, seg_num, seglen, short_size, target_size, img_mean, img_std, name=self.name), label def reader_(): with open(pickle_list) as flist: full_lines = [line.strip() for line in flist] if self.mode == 'train': if (not hasattr(reader_, 'seed')): reader_.seed = 0 random.Random(reader_.seed).shuffle(full_lines) print("reader shuffle seed", reader_.seed) if reader_.seed is not None: reader_.seed += 1 per_node_lines = int( math.ceil(len(full_lines) * 1.0 / self.num_trainers)) total_lines = per_node_lines * self.num_trainers # aligned full_lines so that it can evenly divisible full_lines += full_lines[:(total_lines - len(full_lines))] assert len(full_lines) == total_lines # trainer get own sample lines = full_lines[self.trainer_id:total_lines: self.num_trainers] logger.info("trainerid %d, trainer_count %d" % (self.trainer_id, self.num_trainers)) logger.info( "read images from %d, length: %d, lines length: %d, total: %d" % (self.trainer_id * per_node_lines, per_node_lines, len(lines), len(full_lines))) assert len(lines) == per_node_lines for line in lines: pickle_path = line.strip() yield [pickle_path] if format == 'pkl': decode_func = decode_pickle if format == 'frames': decode_func = decode_frames elif format == 'mp4' or 'avi': decode_func = decode_mp4 else: raise "Not implemented format {}".format(format) mapper = functools.partial( decode_func, mode=mode, seg_num=seg_num, seglen=seglen, short_size=short_size, target_size=target_size, img_mean=img_mean, img_std=img_std) return fluid.io.xmap_readers(mapper, reader_, num_threads, buf_size) def build_dali_reader(self): """ build dali training reader """ def reader_(): with open(self.filelist) as flist: full_lines = [line for line in flist] if self.mode == 'train': if (not hasattr(reader_, 'seed')): reader_.seed = 0 random.Random(reader_.seed).shuffle(full_lines) print("reader shuffle seed", reader_.seed) if reader_.seed is not None: reader_.seed += 1 per_node_lines = int( math.ceil(len(full_lines) * 1.0 / self.num_trainers)) total_lines = per_node_lines * self.num_trainers # aligned full_lines so that it can evenly divisible full_lines += full_lines[:(total_lines - len(full_lines))] assert len(full_lines) == total_lines # trainer get own sample lines = full_lines[self.trainer_id:total_lines: self.num_trainers] assert len(lines) == per_node_lines logger.info("trainerid %d, trainer_count %d" % (self.trainer_id, self.num_trainers)) logger.info( "read images from %d, length: %d, lines length: %d, total: %d" % (self.trainer_id * per_node_lines, per_node_lines, len(lines), len(full_lines))) video_files = '' for item in lines: video_files += item tf = tempfile.NamedTemporaryFile() tf.write(str.encode(video_files)) tf.flush() video_files = tf.name device_id = int(os.getenv('FLAGS_selected_gpus', 0)) print('---------- device id -----------', device_id) if self.mode == 'train': pipe = VideoPipe( batch_size=self.batch_size, num_threads=1, device_id=device_id, file_list=video_files, sequence_length=self.seg_num * self.seglen, seg_num=self.seg_num, seg_length=self.seglen, resize_shorter_scale=self.short_size, crop_target_size=self.target_size, is_training=(self.mode == 'train'), dali_mean=self.dali_mean, dali_std=self.dali_std) else: pipe = VideoTestPipe( batch_size=self.batch_size, num_threads=1, device_id=device_id, file_list=video_files, sequence_length=self.seg_num * self.seglen, seg_num=self.seg_num, seg_length=self.seglen, resize_shorter_scale=self.short_size, crop_target_size=self.target_size, is_training=(self.mode == 'train'), dali_mean=self.dali_mean, dali_std=self.dali_std) logger.info( 'initializing dataset, it will take several minutes if it is too large .... ' ) video_loader = DALIGenericIterator( [pipe], ['image', 'label'], len(lines), dynamic_shape=True, auto_reset=True) return video_loader dali_reader = reader_() def ret_reader(): for data in dali_reader: yield data[0]['image'], data[0]['label'] return ret_reader class VideoPipe(Pipeline): def __init__(self, batch_size, num_threads, device_id, file_list, sequence_length, seg_num, seg_length, resize_shorter_scale, crop_target_size, is_training=False, initial_prefetch_size=10, num_shards=1, shard_id=0, dali_mean=0., dali_std=1.0): super(VideoPipe, self).__init__(batch_size, num_threads, device_id) self.input = ops.VideoReader( device="gpu", file_list=file_list, sequence_length=sequence_length, seg_num=seg_num, seg_length=seg_length, is_training=is_training, num_shards=num_shards, shard_id=shard_id, random_shuffle=is_training, initial_fill=initial_prefetch_size) # the sequece data read by ops.VideoReader is of shape [F, H, W, C] # Because the ops.Resize does not support sequence data, # it will be transposed into [H, W, F, C], # then reshaped to [H, W, FC], and then resized like a 2-D image. self.transpose = ops.Transpose(device="gpu", perm=[1, 2, 0, 3]) self.reshape = ops.Reshape( device="gpu", rel_shape=[1.0, 1.0, -1], layout='HWC') self.resize = ops.Resize( device="gpu", resize_shorter=resize_shorter_scale) # crops and mirror are applied by ops.CropMirrorNormalize. # Normalization will be implemented in paddle due to the difficulty of dimension broadcast, # It is not sure whether dimension broadcast can be implemented correctly by dali, just take the Paddle Op instead. self.pos_rng_x = ops.Uniform(range=(0.0, 1.0)) self.pos_rng_y = ops.Uniform(range=(0.0, 1.0)) self.mirror_generator = ops.Uniform(range=(0.0, 1.0)) self.cast_mirror = ops.Cast(dtype=types.DALIDataType.INT32) self.crop_mirror_norm = ops.CropMirrorNormalize( device="gpu", crop=[crop_target_size, crop_target_size], mean=dali_mean, std=dali_std) self.reshape_back = ops.Reshape( device="gpu", shape=[ seg_num, seg_length * 3, crop_target_size, crop_target_size ], layout='FCHW') self.cast_label = ops.Cast(device="gpu", dtype=types.DALIDataType.INT64) def define_graph(self): output, label = self.input(name="Reader") output = self.transpose(output) output = self.reshape(output) output = self.resize(output) output = output / 255. pos_x = self.pos_rng_x() pos_y = self.pos_rng_y() mirror_flag = self.mirror_generator() mirror_flag = (mirror_flag > 0.5) mirror_flag = self.cast_mirror(mirror_flag) output = self.crop_mirror_norm( output, crop_pos_x=pos_x, crop_pos_y=pos_y, mirror=mirror_flag) output = self.reshape_back(output) label = self.cast_label(label) return output, label class VideoTestPipe(Pipeline): def __init__(self, batch_size, num_threads, device_id, file_list, sequence_length, seg_num, seg_length, resize_shorter_scale, crop_target_size, is_training=False, initial_prefetch_size=10, num_shards=1, shard_id=0, dali_mean=0., dali_std=1.0): super(VideoTestPipe, self).__init__(batch_size, num_threads, device_id) self.input = ops.VideoReader( device="gpu", file_list=file_list, sequence_length=sequence_length, seg_num=seg_num, seg_length=seg_length, is_training=is_training, num_shards=num_shards, shard_id=shard_id, random_shuffle=is_training, initial_fill=initial_prefetch_size) # the sequece data read by ops.VideoReader is of shape [F, H, W, C] # Because the ops.Resize does not support sequence data, # it will be transposed into [H, W, F, C], # then reshaped to [H, W, FC], and then resized like a 2-D image. self.transpose = ops.Transpose(device="gpu", perm=[1, 2, 0, 3]) self.reshape = ops.Reshape( device="gpu", rel_shape=[1.0, 1.0, -1], layout='HWC') self.resize = ops.Resize( device="gpu", resize_shorter=resize_shorter_scale) # crops and mirror are applied by ops.CropMirrorNormalize. # Normalization will be implemented in paddle due to the difficulty of dimension broadcast, # It is not sure whether dimension broadcast can be implemented correctly by dali, just take the Paddle Op instead. self.crop_mirror_norm = ops.CropMirrorNormalize( device="gpu", crop=[crop_target_size, crop_target_size], crop_pos_x=0.5, crop_pos_y=0.5, mirror=0, mean=dali_mean, std=dali_std) self.reshape_back = ops.Reshape( device="gpu", shape=[ seg_num, seg_length * 3, crop_target_size, crop_target_size ], layout='FCHW') self.cast_label = ops.Cast(device="gpu", dtype=types.DALIDataType.INT64) def define_graph(self): output, label = self.input(name="Reader") output = self.transpose(output) output = self.reshape(output) output = self.resize(output) output = output / 255. #output = self.crop(output, crop_pos_x=pos_x, crop_pos_y=pos_y) output = self.crop_mirror_norm(output) output = self.reshape_back(output) label = self.cast_label(label) return output, label def imgs_transform(imgs, mode, seg_num, seglen, short_size, target_size, img_mean, img_std, name=''): imgs = group_scale(imgs, short_size) if mode == 'train': if name == "TSM": imgs = group_multi_scale_crop(imgs, short_size) imgs = group_random_crop(imgs, target_size) imgs = group_random_flip(imgs) else: imgs = group_center_crop(imgs, target_size) np_imgs = (np.array(imgs[0]).astype('float32').transpose( (2, 0, 1))).reshape(1, 3, target_size, target_size) / 255 for i in range(len(imgs) - 1): img = (np.array(imgs[i + 1]).astype('float32').transpose( (2, 0, 1))).reshape(1, 3, target_size, target_size) / 255 np_imgs = np.concatenate((np_imgs, img)) imgs = np_imgs imgs -= img_mean imgs /= img_std imgs = np.reshape(imgs, (seg_num, seglen * 3, target_size, target_size)) return imgs def group_multi_scale_crop(img_group, target_size, scales=None, max_distort=1, fix_crop=True, more_fix_crop=True): scales = scales if scales is not None else [1, .875, .75, .66] input_size = [target_size, target_size] im_size = img_group[0].size # get random crop offset def _sample_crop_size(im_size): image_w, image_h = im_size[0], im_size[1] base_size = min(image_w, image_h) crop_sizes = [int(base_size * x) for x in scales] crop_h = [ input_size[1] if abs(x - input_size[1]) < 3 else x for x in crop_sizes ] crop_w = [ input_size[0] if abs(x - 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) <= max_distort: pairs.append((w, h)) crop_pair = random.choice(pairs) if not fix_crop: w_offset = random.randint(0, image_w - crop_pair[0]) h_offset = random.randint(0, image_h - crop_pair[1]) else: w_step = (image_w - crop_pair[0]) / 4 h_step = (image_h - crop_pair[1]) / 4 ret = list() ret.append((0, 0)) # upper left if w_step != 0: ret.append((4 * w_step, 0)) # upper right if h_step != 0: ret.append((0, 4 * h_step)) # lower left if h_step != 0 and w_step != 0: ret.append((4 * w_step, 4 * h_step)) # lower right if h_step != 0 or w_step != 0: 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 w_offset, h_offset = random.choice(ret) return crop_pair[0], crop_pair[1], w_offset, h_offset crop_w, crop_h, offset_w, offset_h = _sample_crop_size(im_size) crop_img_group = [ img.crop((offset_w, offset_h, offset_w + crop_w, offset_h + crop_h)) for img in img_group ] ret_img_group = [ img.resize((input_size[0], input_size[1]), Image.BILINEAR) for img in crop_img_group ] return ret_img_group def group_random_crop(img_group, target_size): w, h = img_group[0].size th, tw = target_size, target_size assert (w >= target_size) and (h >= target_size), \ "image width({}) and height({}) should be larger than crop size".format( w, h, target_size) out_images = [] x1 = random.randint(0, w - tw) y1 = random.randint(0, h - th) for img in img_group: 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 def group_random_flip(img_group): v = random.random() if v < 0.5: ret = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_group] return ret else: return img_group def group_center_crop(img_group, target_size): img_crop = [] for img in img_group: w, h = img.size th, tw = target_size, target_size assert (w >= target_size) and (h >= target_size), \ "image width({}) and height({}) should be larger than crop size".format( w, h, target_size) x1 = int(round((w - tw) / 2.)) y1 = int(round((h - th) / 2.)) img_crop.append(img.crop((x1, y1, x1 + tw, y1 + th))) return img_crop def group_scale(imgs, target_size): resized_imgs = [] for i in range(len(imgs)): img = imgs[i] w, h = img.size if (w <= h and w == target_size) or (h <= w and h == target_size): resized_imgs.append(img) continue if w < h: ow = target_size oh = int(target_size * 4.0 / 3.0) resized_imgs.append(img.resize((ow, oh), Image.BILINEAR)) else: oh = target_size ow = int(target_size * 4.0 / 3.0) resized_imgs.append(img.resize((ow, oh), Image.BILINEAR)) return resized_imgs def imageloader(buf): if isinstance(buf, str): img = Image.open(StringIO(buf)) else: img = Image.open(BytesIO(buf)) return img.convert('RGB') def video_loader(frames, nsample, seglen, mode): videolen = len(frames) average_dur = int(videolen / nsample) imgs = [] for i in range(nsample): idx = 0 if mode == 'train': if average_dur >= seglen: idx = random.randint(0, average_dur - seglen) idx += i * average_dur elif average_dur >= 1: idx += i * average_dur else: idx = i else: if average_dur >= seglen: idx = (average_dur - seglen) // 2 idx += i * average_dur elif average_dur >= 1: idx += i * average_dur else: idx = i for jj in range(idx, idx + seglen): imgbuf = frames[int(jj % videolen)] img = imageloader(imgbuf) imgs.append(img) return imgs def mp4_loader(filepath, nsample, seglen, mode): cap = cv2.VideoCapture(filepath) videolen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) sampledFrames = [] for i in range(videolen): ret, frame = cap.read() # maybe first frame is empty if ret == False: continue img = frame[:, :, ::-1] sampledFrames.append(img) average_dur = int(len(sampledFrames) / nsample) imgs = [] for i in range(nsample): idx = 0 if mode == 'train': if average_dur >= seglen: idx = random.randint(0, average_dur - seglen) idx += i * average_dur elif average_dur >= 1: idx += i * average_dur else: idx = i else: if average_dur >= seglen: idx = (average_dur - 1) // 2 idx += i * average_dur elif average_dur >= 1: idx += i * average_dur else: idx = i for jj in range(idx, idx + seglen): imgbuf = sampledFrames[int(jj % len(sampledFrames))] img = Image.fromarray(imgbuf, mode='RGB') imgs.append(img) return imgs # the additional function which used to load the frames def load_image(directory, idx): img = cv2.imread(os.path.join(directory, 'img_{:05d}.jpg'.format(idx))) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) return img def frames_loader(recode, nsample, seglen, mode): imgpath, num_frames = recode.path, recode.num_frames average_dur = int(num_frames / nsample) imgs = [] for i in range(nsample): idx = 0 if mode == 'train': if average_dur >= seglen: idx = random.randint(0, average_dur - seglen) idx += i * average_dur elif average_dur >= 1: idx += i * average_dur else: idx = i else: if average_dur >= seglen: idx = (average_dur - 1) // 2 idx += i * average_dur elif average_dur >= 1: idx += i * average_dur else: idx = i for jj in range(idx, idx + seglen): img = load_image(imgpath, jj + 1) img = Image.fromarray(img, mode='RGB') imgs.append(img) return imgs