def _aggregate_per_chunk_results(self, algorithm_type): lr_video_path = os.path.join(self.dataset_dir, 'video', self.lr_video_name) lr_video_profile = get_video_profile(lr_video_path) num_chunks = int(math.ceil(lr_video_profile['duration'] / (self.gop / lr_video_profile['frame_rate']))) start_idx = 0 end_idx = num_chunks - 1 log_dir = os.path.join(self.dataset_dir, 'log', self.lr_video_name, self.model.name) cache_profile_dir = os.path.join(self.dataset_dir, 'profile', self.lr_video_name, self.model.name) log_name = os.path.join('quality_{}.txt'.format(algorithm_type)) cache_profile_name = os.path.join('{}.profile'.format(algorithm_type)) #log log_path = os.path.join(log_dir, log_name) with open(log_path, 'w') as f0: #iterate over chunks for chunk_idx in range(start_idx, end_idx + 1): chunk_log_dir = os.path.join(log_dir, 'chunk{:04d}'.format(chunk_idx)) chunk_log_path= os.path.join(chunk_log_dir, log_name) with open(chunk_log_path, 'r') as f1: q_lines = f1.readlines() f0.write('{}\t{}\n'.format(chunk_idx, q_lines[-1].strip())) #cache profile cache_profile_path = os.path.join(cache_profile_dir, cache_profile_name) cache_data = b'' with open(cache_profile_path, 'wb') as f0: for chunk_idx in range(start_idx, end_idx + 1): chunk_cache_profile_path = os.path.join(cache_profile_dir, 'chunk{:04d}'.format(chunk_idx), cache_profile_name) with open(chunk_cache_profile_path, 'rb') as f1: f0.write(f1.read()) #log (bilinear, sr) log_path = os.path.join(self.dataset_dir, 'log', self.lr_video_name, 'quality.txt') with open(log_path, 'w') as f0: #iterate over chunks for chunk_idx in range(start_idx, end_idx + 1): quality = [] chunk_log_path = os.path.join(self.dataset_dir, 'log', self.lr_video_name, 'chunk{:04d}'.format(chunk_idx), 'quality.txt') with open(chunk_log_path, 'r') as f1: lines = f1.readlines() for line in lines: line = line.strip() quality.append(float(line.split('\t')[1])) f0.write('{}\t{:.4f}\n'.format(chunk_idx, np.average(quality))) log_path = os.path.join(self.dataset_dir, 'log', self.lr_video_name, self.model.name, 'quality.txt') with open(log_path, 'w') as f0: #iterate over chunks for chunk_idx in range(start_idx, end_idx + 1): quality = [] chunk_log_path = os.path.join(self.dataset_dir, 'log', self.lr_video_name, self.model.name, 'chunk{:04d}'.format(chunk_idx), 'quality.txt') with open(chunk_log_path, 'r') as f1: lines = f1.readlines() for line in lines: line = line.strip() quality.append(float(line.split('\t')[1])) f0.write('{}\t{:.4f}\n'.format(chunk_idx, np.average(quality)))
def select_anchor_point_set(self, algorithm_type, chunk_idx=None, max_nemo_num_anchor_points=None): if chunk_idx is not None: if algorithm_type == 'nemo': self._select_anchor_point_set_nemo(chunk_idx) elif algorithm_type == 'uniform': self._select_anchor_point_set_uniform(chunk_idx) elif algorithm_type == 'random': self._select_anchor_point_set_random(chunk_idx) else: lr_video_path = os.path.join(self.dataset_dir, 'video', self.lr_video_name) lr_video_profile = get_video_profile(lr_video_path) num_chunks = int(math.ceil(lr_video_profile['duration'] / (self.gop / lr_video_profile['frame_rate']))) for i in range(num_chunks): if algorithm_type == 'nemo': self._select_anchor_point_set_nemo(i) elif algorithm_type == 'uniform': self._select_anchor_point_set_uniform(i) elif algorithm_type == 'random': self._select_anchor_point_set_random(i)
#path parser.add_argument('--ffmpeg_path', type=str, default='/usr/bin/ffmpeg') parser.add_argument('--input_video_path', type=str, required=True) parser.add_argument('--output_video_dir', type=str, required=True) #video parser.add_argument('--mode', type=str, required=True) parser.add_argument('--bitrate', type=int, default=None) parser.add_argument('--output_width', type=int, default=None) parser.add_argument('--output_height', type=int, default=None) parser.add_argument('--gop', type=int, default=120) parser.add_argument('--start', type=int, default=None) parser.add_argument('--duration', type=int, default=None) args = parser.parse_args() print(args.input_video_path) input_video_height = get_video_profile(args.input_video_path)['height'] if args.mode == 'cut_and_resize_and_encode': assert(args.start is not None and args.duration is not None) enc = LibvpxEncoder(args.output_video_dir, args.input_video_path, input_video_height, args.start, args.duration, args.ffmpeg_path) enc.cut_and_resize_and_encode(args.output_width, args.output_height, args.bitrate, args.gop) elif args.mode == 'resize_and_encode': assert(args.bitrate is not None and args.output_height is not None and args.output_width is not None) enc = LibvpxEncoder(args.output_video_dir, args.input_video_path, input_video_height, args.start, args.duration, args.ffmpeg_path) enc.resize_and_encode(args.output_width, args.output_height, args.bitrate, args.gop) else: raise ValueError('Unsupported mode')
choices=['nemo', 'uniform', 'random'], required=True) args = parser.parse_args() if args.vpxdec_path is None: args.vpxdec_path = os.path.join(os.environ['NEMO_CODE_ROOT'], 'third_party', 'libvpx', 'bin', 'vpxdec_nemo_ver2_x86') assert (os.path.exists(args.vpxdec_path)) # setup dataset_dir = os.path.join(args.data_dir, args.content) lr_video_path = os.path.join(dataset_dir, 'video', args.lr_video_name) hr_video_path = os.path.join(dataset_dir, 'video', args.hr_video_name) lr_video_profile = get_video_profile(lr_video_path) hr_video_profile = get_video_profile(hr_video_path) scale = args.output_height // lr_video_profile['height'] model = build_model(args.model_type, args.num_blocks, args.num_filters, scale, args.upsample_type) # prepare metadata metadata_file = os.path.join(dataset_dir, 'log', args.lr_video_name, model.name, 'metadata.txt') cache_profile_name = '{}_{}'.format(args.algorithm, args.quality_margin) if not os.path.exists(metadata_file): offline_cache_metadata(args.vpxdec_path, dataset_dir, args.lr_video_name, model.name, cache_profile_name, args.output_width, args.output_height)
def _select_anchor_point_set_nemo(self, chunk_idx): postfix = 'chunk{:04d}'.format(chunk_idx) cache_profile_dir = os.path.join(self.dataset_dir, 'profile', self.lr_video_name, self.model.name, postfix) log_dir = os.path.join(self.dataset_dir, 'log', self.lr_video_name, self.model.name, postfix) os.makedirs(log_dir, exist_ok=True) os.makedirs(cache_profile_dir, exist_ok=True) algorithm_type = 'nemo_{}'.format(self.quality_margin) ###########step 1: analyze anchor points########## #calculate num_skipped_frames and num_decoded frames start_time = time.time() lr_video_path = os.path.join(self.dataset_dir, 'video', self.lr_video_name) lr_video_profile = get_video_profile(lr_video_path) num_total_frames = int(round(lr_video_profile['frame_rate'], 3) * round(lr_video_profile['duration'])) num_left_frames = num_total_frames - chunk_idx * self.gop assert(num_total_frames == math.floor(num_total_frames)) num_skipped_frames = chunk_idx * self.gop num_decoded_frames = self.gop if num_left_frames >= self.gop else num_left_frames #save low-resolution, super-resoluted, high-resolution frames to local storage libvpx.save_rgb_frame(self.vpxdec_path, self.dataset_dir, self.lr_video_name, skip=num_skipped_frames, limit=num_decoded_frames, postfix=postfix) libvpx.save_yuv_frame(self.vpxdec_path, self.dataset_dir, self.hr_video_name, self.output_width, self.output_height, num_skipped_frames, num_decoded_frames, postfix) libvpx.setup_sr_frame(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.model, postfix) #measure bilinear, per-frame super-resolution quality quality_bilinear = libvpx.bilinear_quality(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.hr_video_name, self.output_width, self.output_height, num_skipped_frames, num_decoded_frames, postfix) quality_dnn = libvpx.offline_dnn_quality(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.hr_video_name, self.model.name, \ self.output_width, self.output_height, num_skipped_frames, num_decoded_frames, postfix) end_time = time.time() print('{} video chunk: (Step1-profile bilinear, dnn quality) {}sec'.format(chunk_idx, end_time - start_time)) #create multiple processes for parallel quality measurements start_time = time.time() q0 = mp.Queue() q1 = mp.Queue() decoders = [mp.Process(target=libvpx.offline_cache_quality_mt, args=(q0, q1, self.vpxdec_path, self.dataset_dir, \ self.lr_video_name, self.hr_video_name, self.model.name, self.output_width, self.output_height)) for i in range(self.num_decoders)] for decoder in decoders: decoder.start() #select a single anchor point and measure the resulting quality single_anchor_point_sets = [] frames = libvpx.load_frame_index(self.dataset_dir, self.lr_video_name, postfix) for idx, frame in enumerate(frames): anchor_point_set = libvpx.AnchorPointSet.create(frames, cache_profile_dir, frame.name) anchor_point_set.add_anchor_point(frame) anchor_point_set.save_cache_profile() q0.put((anchor_point_set.get_cache_profile_name(), num_skipped_frames, num_decoded_frames, postfix, idx)) single_anchor_point_sets.append(anchor_point_set) for frame in frames: item = q1.get() idx = item[0] quality = item[1] single_anchor_point_sets[idx].set_measured_quality(quality) single_anchor_point_sets[idx].remove_cache_profile() #remove multiple processes for decoder in decoders: q0.put('end') for decoder in decoders: decoder.join() end_time = time.time() print('{} video chunk: (Step1-profile anchor point quality) {}sec'.format(chunk_idx, end_time - start_time)) ###########step 2: order anchor points########## start_time = time.time() multiple_anchor_point_sets = [] anchor_point_set = None FAST_anchor_point_set = single_anchor_point_sets[0] while len(single_anchor_point_sets) > 0: anchor_point_idx, estimated_quality = self._select_anchor_point(anchor_point_set, single_anchor_point_sets) selected_anchor_point = single_anchor_point_sets.pop(anchor_point_idx) if len(multiple_anchor_point_sets) == 0: anchor_point_set = libvpx.AnchorPointSet.load(selected_anchor_point, cache_profile_dir, '{}_{}'.format(algorithm_type, 1)) anchor_point_set.set_estimated_quality(selected_anchor_point.measured_quality) else: anchor_point_set = libvpx.AnchorPointSet.load(multiple_anchor_point_sets[-1], cache_profile_dir, '{}_{}'.format(algorithm_type, multiple_anchor_point_sets[-1].get_num_anchor_points() + 1)) anchor_point_set.add_anchor_point(selected_anchor_point.anchor_points[0]) anchor_point_set.set_estimated_quality(estimated_quality) multiple_anchor_point_sets.append(anchor_point_set) end_time = time.time() print('{} video chunk: (Step2) {}sec'.format(chunk_idx, end_time - start_time)) ###########step 3: select anchor points########## start_time = time.time() log_path0 = os.path.join(log_dir, 'quality_{}.txt'.format(algorithm_type)) log_path1 = os.path.join(log_dir, 'quality_{}_8.txt'.format(algorithm_type)) log_path2 = os.path.join(log_dir, 'quality_{}_16.txt'.format(algorithm_type)) log_path3 = os.path.join(log_dir, 'quality_fast.txt') with open(log_path0, 'w') as f0, open(log_path1, 'w') as f1, open(log_path2, 'w') as f2, open(log_path3, 'w') as f3: for idx, anchor_point_set in enumerate(multiple_anchor_point_sets): #log quality anchor_point_set.save_cache_profile() quality_cache = libvpx.offline_cache_quality(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.hr_video_name, \ self.model.name, anchor_point_set.get_cache_profile_name(), self.output_width, self.output_height, \ num_skipped_frames, num_decoded_frames, postfix) anchor_point_set.remove_cache_profile() quality_diff = np.asarray(quality_dnn) - np.asarray(quality_cache) quality_log = '{}\t{}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}\n'.format(anchor_point_set.get_num_anchor_points(), len(frames), \ np.average(quality_cache), np.average(quality_dnn), np.average(quality_bilinear), np.average(anchor_point_set.estimated_quality)) f0.write(quality_log) if idx < 8: f1.write(quality_log) if idx < 16: f2.write(quality_log) if idx == 0: f3.write(quality_log) #terminate if np.average(quality_diff) <= self.quality_margin: #case 1: does not restrict #anchor points anchor_point_set.set_cache_profile_name(algorithm_type) anchor_point_set.save_cache_profile() #case 2: limit #anchor points to 8 if anchor_point_set.get_num_anchor_points() > 8: anchor_point_set_ = multiple_anchor_point_sets[7] else: anchor_point_set_ = anchor_point_set anchor_point_set_.set_cache_profile_name('{}_8'.format(algorithm_type)) anchor_point_set_.save_cache_profile() #case 3: limit #anchor points to 16 if anchor_point_set.get_num_anchor_points() > 16: anchor_point_set_ = multiple_anchor_point_sets[15] else: anchor_point_set_ = anchor_point_set anchor_point_set_.set_cache_profile_name('{}_16'.format(algorithm_type)) anchor_point_set_.save_cache_profile() #case 4: FAST anchor_point_set_ = FAST_anchor_point_set anchor_point_set_.set_cache_profile_name('fast'.format(algorithm_type)) anchor_point_set_.save_cache_profile() break end_time = time.time() print('{} video chunk: (Step3) {}sec'.format(chunk_idx, end_time - start_time)) #remove images lr_image_dir = os.path.join(self.dataset_dir, 'image', self.lr_video_name, postfix) hr_image_dir = os.path.join(self.dataset_dir, 'image', self.hr_video_name, postfix) sr_image_dir = os.path.join(self.dataset_dir, 'image', self.lr_video_name, self.model.name, postfix) shutil.rmtree(lr_image_dir, ignore_errors=True) shutil.rmtree(hr_image_dir, ignore_errors=True) shutil.rmtree(sr_image_dir, ignore_errors=True)
def _select_anchor_point_set_random(self, chunk_idx=None): postfix = 'chunk{:04d}'.format(chunk_idx) cache_profile_dir = os.path.join(self.dataset_dir, 'profile', self.lr_video_name, self.model.name, postfix) log_dir = os.path.join(self.dataset_dir, 'log', self.lr_video_name, self.model.name, postfix) os.makedirs(log_dir, exist_ok=True) os.makedirs(cache_profile_dir, exist_ok=True) algorithm_type = 'random_{}'.format(self.quality_margin) ###########step 1: measure bilinear, dnn quality########## #calculate num_skipped_frames and num_decoded frames start_time = time.time() lr_video_path = os.path.join(self.dataset_dir, 'video', self.lr_video_name) lr_video_profile = get_video_profile(lr_video_path) num_total_frames = int(round(lr_video_profile['frame_rate'], 3) * round(lr_video_profile['duration'])) num_left_frames = num_total_frames - chunk_idx * self.gop assert(num_total_frames == math.floor(num_total_frames)) num_skipped_frames = chunk_idx * self.gop num_decoded_frames = self.gop if num_left_frames >= self.gop else num_left_frames #save low-resolution, super-resoluted, high-resolution frames to local storage libvpx.save_rgb_frame(self.vpxdec_path, self.dataset_dir, self.lr_video_name, skip=num_skipped_frames, limit=num_decoded_frames, postfix=postfix) libvpx.save_yuv_frame(self.vpxdec_path, self.dataset_dir, self.hr_video_name, self.output_width, self.output_height, num_skipped_frames, num_decoded_frames, postfix) libvpx.setup_sr_frame(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.model, postfix) #measure bilinear, per-frame super-resolution quality quality_bilinear = libvpx.bilinear_quality(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.hr_video_name, self.output_width, self.output_height, num_skipped_frames, num_decoded_frames, postfix) quality_dnn = libvpx.offline_dnn_quality(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.hr_video_name, self.model.name, \ self.output_width, self.output_height, num_skipped_frames, num_decoded_frames, postfix) end_time = time.time() print('{} video chunk: (Step1-profile bilinear, dnn quality) {}sec'.format(chunk_idx, end_time - start_time)) ###########step 2: select anchor points########## start_time = time.time() frames = libvpx.load_frame_index(self.dataset_dir, self.lr_video_name, postfix) log_path = os.path.join(log_dir, 'quality_{}.txt'.format(algorithm_type)) with open(log_path, 'w') as f: for i in range(len(frames)): #select anchor point uniformly num_anchor_points = i + 1 anchor_point_set = libvpx.AnchorPointSet.create(frames, cache_profile_dir, '{}_{}'.format(algorithm_type, num_anchor_points)) random_frames = random.sample(frames, num_anchor_points) for frame in random_frames: anchor_point_set.add_anchor_point(frame) #measure the quality anchor_point_set.save_cache_profile() quality_cache = libvpx.offline_cache_quality(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.hr_video_name, \ self.model.name, anchor_point_set.get_cache_profile_name(), self.output_width, self.output_height, \ num_skipped_frames, num_decoded_frames, postfix) anchor_point_set.remove_cache_profile() quality_diff = np.asarray(quality_dnn) - np.asarray(quality_cache) quality_log = '{}\t{}\t{:.4f}\t{:.4f}\t{:.4f}\n'.format(anchor_point_set.get_num_anchor_points(), len(frames), \ np.average(quality_cache), np.average(quality_dnn), np.average(quality_bilinear)) f.write(quality_log) #terminate if np.average(quality_diff) <= self.quality_margin: anchor_point_set.set_cache_profile_name(algorithm_type) anchor_point_set.save_cache_profile() libvpx.offline_cache_quality(self.vpxdec_path, self.dataset_dir, self.lr_video_name, self.hr_video_name, \ self.model.name, anchor_point_set.get_cache_profile_name(), self.output_width, self.output_height, \ num_skipped_frames, num_decoded_frames, postfix) break end_time = time.time() print('{} video chunk: (Step2) {}sec'.format(chunk_idx, end_time - start_time)) #remove images lr_image_dir = os.path.join(self.dataset_dir, 'image', self.lr_video_name, postfix) hr_image_dir = os.path.join(self.dataset_dir, 'image', self.hr_video_name, postfix) sr_image_dir = os.path.join(self.dataset_dir, 'image', self.lr_video_name, self.model.name, postfix) shutil.rmtree(lr_image_dir, ignore_errors=True) shutil.rmtree(hr_image_dir, ignore_errors=True) shutil.rmtree(sr_image_dir, ignore_errors=True)
#directory, path parser.add_argument('--dataset_dir', type=str, required=True) parser.add_argument('--content', type=str, required=True) parser.add_argument('--video_name', type=str, required=True) #architecture parser.add_argument('--model_type', type=str, default='nemo_s') parser.add_argument('--num_filters', type=int, required=True) parser.add_argument('--num_blocks', type=int, required=True) parser.add_argument('--upsample_type', type=str, default='deconv') parser.add_argument('--scale', type=int, required=True) args = parser.parse_args() #scale, dnn video_path = os.path.join(args.dataset_dir, args.content, 'video', args.video_name) video_profile = get_video_profile(video_path) with tf.Graph().as_default(), tf.Session() as sess: init = tf.global_variables_initializer() sess.run(init) model = build_model(args.model_type, args.num_blocks, args.num_filters, args.scale, args.upsample_type) log_dir = os.path.join(args.dataset_dir, args.content, 'log', args.video_name, model.name) os.makedirs(log_dir, exist_ok=True) summary_writer = tf.summary.FileWriter(log_dir, sess.graph)