diff --git a/applications/DAIN/my_args.py b/applications/DAIN/my_args.py index 7a5df29b04f6c4d4aeb3b61dcb19a10c5c803b88..448c3c2215c4a5fcba61bcb79ca242d3a3db1b18 100644 --- a/applications/DAIN/my_args.py +++ b/applications/DAIN/my_args.py @@ -90,4 +90,7 @@ parser.add_argument('--use_cuda', type=bool, help='use cuda or not') parser.add_argument('--use_cudnn', default=1, type=int, help='use cudnn or not') - +parser.add_argument('--remove_duplicates', + default=True, + type=bool, + help='remove duplicate frames or not') diff --git a/applications/DAIN/predict.py b/applications/DAIN/predict.py index f43a794150d52e452fdb7d5c3aadcc93cbbda3da..39290433b2b0cfbc2bfcce18d800740987c1d131 100644 --- a/applications/DAIN/predict.py +++ b/applications/DAIN/predict.py @@ -80,7 +80,8 @@ class VideoFrameInterp(object): video_path, use_gpu=True, key_frame_thread=0., - output_path='output'): + output_path='output', + remove_duplicates=True): self.video_path = video_path self.output_path = os.path.join(output_path, 'DAIN') if model_path is None: @@ -138,6 +139,8 @@ class VideoFrameInterp(object): end = time.time() frames = sorted(glob.glob(os.path.join(out_path, '*.png'))) + if remove_duplicates: + frames = remove_duplicates(out_path) img = imread(frames[0]) @@ -199,58 +202,51 @@ class VideoFrameInterp(object): X0 = img_first.astype('float32').transpose((2, 0, 1)) / 255 X1 = img_second.astype('float32').transpose((2, 0, 1)) / 255 - if key_frame: - y_ = [ - np.transpose(255.0 * X0.clip(0, 1.0), (1, 2, 0)) - for i in range(num_frames) - ] - else: - assert (X0.shape[1] == X1.shape[1]) - assert (X0.shape[2] == X1.shape[2]) - - X0 = np.pad(X0, ((0,0), (padding_top, padding_bottom), \ - (padding_left, padding_right)), mode='edge') - X1 = np.pad(X1, ((0,0), (padding_top, padding_bottom), \ - (padding_left, padding_right)), mode='edge') - - X0 = np.expand_dims(X0, axis=0) - X1 = np.expand_dims(X1, axis=0) - - X0 = np.expand_dims(X0, axis=0) - X1 = np.expand_dims(X1, axis=0) - - X = np.concatenate((X0, X1), axis=0) - - proc_end = time.time() - o = self.exe.run(self.program, - fetch_list=self.fetch_targets, - feed={"image": X}) - - y_ = o[0] - - proc_timer.update(time.time() - proc_end) - tot_timer.update(time.time() - end) - end = time.time() - - y_ = [ - np.transpose( - 255.0 * item.clip( - 0, 1.0)[0, :, - padding_top:padding_top + int_height, - padding_left:padding_left + int_width], - (1, 2, 0)) for item in y_ - ] - time_offsets = [ - kk * timestep for kk in range(1, 1 + num_frames, 1) - ] - - count = 1 - for item, time_offset in zip(y_, time_offsets): - out_dir = os.path.join( - frame_path_interpolated, vidname, - "{:0>6d}_{:0>4d}.png".format(i, count)) - count = count + 1 - imsave(out_dir, np.round(item).astype(np.uint8)) + assert (X0.shape[1] == X1.shape[1]) + assert (X0.shape[2] == X1.shape[2]) + + X0 = np.pad(X0, ((0,0), (padding_top, padding_bottom), \ + (padding_left, padding_right)), mode='edge') + X1 = np.pad(X1, ((0,0), (padding_top, padding_bottom), \ + (padding_left, padding_right)), mode='edge') + + X0 = np.expand_dims(X0, axis=0) + X1 = np.expand_dims(X1, axis=0) + + X0 = np.expand_dims(X0, axis=0) + X1 = np.expand_dims(X1, axis=0) + + X = np.concatenate((X0, X1), axis=0) + + proc_end = time.time() + o = self.exe.run(self.program, + fetch_list=self.fetch_targets, + feed={"image": X}) + + y_ = o[0] + + proc_timer.update(time.time() - proc_end) + tot_timer.update(time.time() - end) + end = time.time() + + y_ = [ + np.transpose( + 255.0 * item.clip( + 0, 1.0)[0, :, padding_top:padding_top + int_height, + padding_left:padding_left + int_width], + (1, 2, 0)) for item in y_ + ] + time_offsets = [ + kk * timestep for kk in range(1, 1 + num_frames, 1) + ] + + count = 1 + for item, time_offset in zip(y_, time_offsets): + out_dir = os.path.join( + frame_path_interpolated, vidname, + "{:0>6d}_{:0>4d}.png".format(i, count)) + count = count + 1 + imsave(out_dir, np.round(item).astype(np.uint8)) num_frames = int(1.0 / timestep) - 1 @@ -266,14 +262,16 @@ class VideoFrameInterp(object): vidname + '.mp4') if os.path.exists(video_pattern_output): os.remove(video_pattern_output) - frames2video(frame_pattern_combined, video_pattern_output, - r2) + frames2video(frame_pattern_combined, video_pattern_output, r2) return frame_pattern_combined, video_pattern_output if __name__ == '__main__': args = parser.parse_args() - predictor = VideoFrameInterp(args.time_step, args.saved_model, - args.video_path, args.output_path) + predictor = VideoFrameInterp(args.time_step, + args.saved_model, + args.video_path, + args.output_path, + remove_duplicates=args.remove_duplicates) predictor.run() diff --git a/applications/DAIN/util.py b/applications/DAIN/util.py index 24ea2741517660581c12d8b174e3e8af03ae9a8e..dc343ff20668f0b9db3a426a3789943ea4e28cb5 100644 --- a/applications/DAIN/util.py +++ b/applications/DAIN/util.py @@ -1,6 +1,7 @@ import os, sys import glob import shutil +import cv2 class AverageMeter(object): @@ -44,3 +45,34 @@ def combine_frames(input, interpolated, combined, num_frames): except Exception as e: print(e) print(len(frames2), num_frames, i, k, i * num_frames + k) + + +def remove_duplicates(paths): + def dhash(image, hash_size=8): + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + resized = cv2.resize(gray, (hash_size + 1, hash_size)) + diff = resized[:, 1:] > resized[:, :-1] + return sum([2**i for (i, v) in enumerate(diff.flatten()) if v]) + + hashes = {} + image_paths = sorted(glob.glob(os.path.join(paths, '*.png'))) + for image_path in image_paths: + image = cv2.imread(image_path) + h = dhash(image) + p = hashes.get(h, []) + p.append(image_path) + hashes[h] = p + + for (h, hashed_paths) in hashes.items(): + if len(hashed_paths) > 1: + for p in hashed_paths[1:]: + os.remove(p) + + frames = sorted(glob.glob(os.path.join(paths, '*.png'))) + for fid, frame in enumerate(frames): + new_name = '{:08d}'.format(fid) + '.png' + new_name = os.path.join(paths, new_name) + os.rename(frame, new_name) + + frames = sorted(glob.glob(os.path.join(paths, '*.png'))) + return frames