From 45165914ba652f1b617e86a3e05059481ba3122f Mon Sep 17 00:00:00 2001 From: kinghuin Date: Thu, 14 May 2020 15:13:35 +0800 Subject: [PATCH] accelerate video reader --- .../videotag_tsn_lstm/module.py | 14 +- .../resource/reader/kinetics_reader.py | 267 ++---------------- 2 files changed, 28 insertions(+), 253 deletions(-) diff --git a/hub_module/modules/video/classification/videotag_tsn_lstm/module.py b/hub_module/modules/video/classification/videotag_tsn_lstm/module.py index dc70f8f2..5390e732 100644 --- a/hub_module/modules/video/classification/videotag_tsn_lstm/module.py +++ b/hub_module/modules/video/classification/videotag_tsn_lstm/module.py @@ -88,12 +88,9 @@ class VideoTag(hub.Module): extractor_model.load_test_weights(exe, args.extractor_weights, extractor_main_prog) - # get reader and metrics - extractor_reader = get_reader("TSN", 'infer', - extractor_infer_config) extractor_feeder = fluid.DataFeeder( place=place, feed_list=extractor_feeds) - return extractor_reader, extractor_main_prog, extractor_fetch_list, extractor_feeder, extractor_scope + return extractor_main_prog, extractor_fetch_list, extractor_feeder, extractor_scope def _predictor(self, args, exe, place): predictor_scope = fluid.Scope() @@ -169,15 +166,20 @@ class VideoTag(hub.Module): self.place = fluid.CUDAPlace( 0) if args.use_gpu else fluid.CPUPlace() self.exe = fluid.Executor(self.place) - self.extractor_reader, self.extractor_main_prog, self.extractor_fetch_list, self.extractor_feeder, self.extractor_scope = self._extractor( + self.extractor_main_prog, self.extractor_fetch_list, self.extractor_feeder, self.extractor_scope = self._extractor( args, self.exe, self.place) self.predictor_main_prog, self.predictor_fetch_list, self.predictor_feeder, self.predictor_scope = self._predictor( args, self.exe, self.place) self._has_load = True + extractor_config = parse_config(args.extractor_config) + extractor_infer_config = merge_configs(extractor_config, 'infer', + vars(args)) + extractor_reader = get_reader("TSN", 'infer', extractor_infer_config) feature_list = [] file_list = [] - for idx, data in enumerate(self.extractor_reader()): + + for idx, data in enumerate(extractor_reader()): file_id = [item[-1] for item in data] feed_data = [item[:-1] for item in data] feature_out = self.exe.run( diff --git a/hub_module/modules/video/classification/videotag_tsn_lstm/resource/reader/kinetics_reader.py b/hub_module/modules/video/classification/videotag_tsn_lstm/resource/reader/kinetics_reader.py index ccf59ea2..2c42da5c 100644 --- a/hub_module/modules/video/classification/videotag_tsn_lstm/resource/reader/kinetics_reader.py +++ b/hub_module/modules/video/classification/videotag_tsn_lstm/resource/reader/kinetics_reader.py @@ -87,38 +87,15 @@ class KineticsReader(DataReader): def _batch_reader(): batch_out = [] for imgs, label in _reader(): - #for imgs in _reader(): if imgs is None: continue batch_out.append((imgs, label)) - #batch_out.append((imgs,)) 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(video_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, @@ -149,37 +126,7 @@ class KineticsReader(DataReader): return imgs_transform(imgs, mode, seg_num, seglen, \ short_size, target_size, img_mean, img_std, name = self.name), mp4_path - 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 reader(): - # with open(pickle_list) as flist: - # lines = [line.strip() for line in flist] lines = [line.strip() for line in pickle_list] if shuffle: random.shuffle(lines) @@ -187,15 +134,8 @@ class KineticsReader(DataReader): pickle_path = line.strip() yield [pickle_path] - if format == 'pkl': - decode_func = decode_pickle - elif format == 'mp4': - decode_func = decode_mp4 - else: - raise "Not implemented format {}".format(format) - mapper = functools.partial( - decode_func, + decode_mp4, mode=mode, seg_num=seg_num, seglen=seglen, @@ -218,142 +158,26 @@ def imgs_transform(imgs, 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) + np_imgs = np.array([np.array(img).astype('float32') for img in imgs]) #dhwc + np_imgs = group_center_crop(np_imgs, target_size) + np_imgs = np_imgs.transpose(0, 3, 1, 2) / 255 #dchw + np_imgs -= img_mean + np_imgs /= img_std - return crop_pair[0], crop_pair[1], w_offset, h_offset + return np_imgs - 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 +def group_center_crop(np_imgs, target_size): + d, h, w, c = np_imgs.shape 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) + "image width({}) and height({}) should be larger than crop size".format(w, h, target_size) - 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))) + h_off = int(round((h - th) / 2.)) + w_off = int(round((w - tw) / 2.)) + img_crop = np_imgs[:, h_off:h_off + target_size, w_off:w_off + + target_size, :] return img_crop @@ -378,47 +202,6 @@ def group_scale(imgs, target_size): 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)) @@ -434,26 +217,16 @@ def mp4_loader(filepath, nsample, seglen, mode): 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 + if average_dur >= seglen: + idx = (average_dur - 1) // 2 + idx += i * average_dur + elif average_dur >= 1: + idx += i * average_dur 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 + 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 -- GitLab