diff --git a/inference_video_parallel.py b/inference_video_parallel.py index 1c92797..eaa0d7f 100644 --- a/inference_video_parallel.py +++ b/inference_video_parallel.py @@ -18,10 +18,10 @@ parser.add_argument('--skip', dest='skip', action='store_true', help='whether to parser.add_argument('--fps', dest='fps', type=int, default=None) parser.add_argument('--png', dest='png', action='store_true', help='whether to output png format outputs') parser.add_argument('--ext', dest='ext', type=str, default='mp4', help='output video extension') -parser.add_argument('--times', dest='times', type=int, default=1) +parser.add_argument('--exp', dest='exp', type=int, default=1, help='interpolation exponent (base 2)') args = parser.parse_args() -assert (args.times == 1 or args.times == 2) -args.times = 2 ** args.times +assert (args.exp in [1, 2, 3]) +args.times = 2 ** args.exp from model.RIFE import Model model = Model() @@ -36,47 +36,58 @@ h, w, _ = frame.shape if args.fps is None: args.fps = fps * args.times fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') +video_path_wo_ext, ext = os.path.splitext(args.video) if args.png: if not os.path.exists('output'): os.mkdir('output') + vid_out = None else: - video_path_wo_ext, ext = os.path.splitext(args.video) - output = cv2.VideoWriter('{}_{}X_{}fps.{}'.format(video_path_wo_ext, args.times, int(np.round(args.fps)), args.ext), fourcc, args.fps, (w, h)) + vid_out = cv2.VideoWriter('{}_{}X_{}fps.{}'.format(video_path_wo_ext, args.times, int(np.round(args.fps)), args.ext), fourcc, args.fps, (w, h)) cnt = 0 skip_frame = 1 -def writeframe(I0, mid0, mid1, mid2, I1, p): - global cnt, skip_frame, args - for i in range(I0.shape[0]): + + +def write_frame(vid_out, i0, infs, i1, p, user_args): + global skip_frame, cnt + + for i in range(i0.shape[0]): + # Result was not good enough to write, use previous frames. if p[i] > 0.2: - if args.times == 4: - mid0[i] = I0[i] - mid1[i] = I0[i] - if args.times == 4: - mid2[i] = I1[i] - if p[i] < 2e-3 and args.skip: + if user_args.exp > 1: + infs = [i0[i] for _ in range(len(infs) - 1)] + infs[-1] = i1[-1] + else: + infs = [i0[i] for _ in range(len(infs))] + + # Result was too similar to previous frame, skip if given. + if p[i] < 2e-3 and user_args.skip: if skip_frame % 100 == 0: - print("Warning: Your video has {} static frames, skipping them may change the duration of the generated video.".format(skip_frame)) + print("Warning: Your video has {} static frames, " + "skipping them may change the duration of the generated video.".format(skip_frame)) skip_frame += 1 continue - if args.png: - cv2.imwrite('output/{:0>7d}.png'.format(cnt), I0[i]) + + # Write results. + if user_args.png: + cv2.imwrite('output/{:0>7d}.png'.format(cnt), i0[i]) cnt += 1 - if args.times == 4: - cv2.imwrite('output/{:0>7d}.png'.format(cnt), mid0[i]) - cnt += 1 - cv2.imwrite('output/{:0>7d}.png'.format(cnt), mid1[i]) - cnt += 1 - if args.times == 4: - cv2.imwrite('output/{:0>7d}.png'.format(cnt), mid2[i]) + for inf in infs: + cv2.imwrite('output/{:0>7d}.png'.format(cnt), inf[i]) cnt += 1 else: - output.write(I0[i]) - if args.times == 4: - output.write(mid0[i]) - output.write(mid1[i]) - if args.times == 4: - output.write(mid2[i]) + vid_out.write(i0[i]) + for inf in infs: + vid_out.write(inf[i]) + + +def make_inference(model, I0, I1, exp): + middle = model.inference(I0, I1) + if exp == 1: + return [middle] + first_half = make_inference(model, I0, middle, exp=exp - 1) + second_half = make_inference(model, middle, I1, exp=exp - 1) + return [*first_half, middle, *second_half] ph = ((h - 1) // 32 + 1) * 32 @@ -97,19 +108,14 @@ while success: - F.interpolate(I1, (16, 16), mode='bilinear', align_corners=False)).abs() I0 = F.pad(I0, padding) I1 = F.pad(I1, padding) - mid1 = model.inference(I0, I1) - if args.times == 4: - mid0 = model.inference(I0, mid1) - mid2 = model.inference(mid1, I1) + inferences = make_inference(model, I0, I1, exp=args.exp) + I0 = ((I0[:, :, :h, :w] * 255.).cpu().detach().numpy().transpose(0, 2, 3, 1)).astype('uint8') I1 = ((I1[:, :, :h, :w] * 255.).cpu().detach().numpy().transpose(0, 2, 3, 1)).astype('uint8') - mid1 = ((mid1[:, :, :h, :w] * 255.).cpu().detach().numpy().transpose(0, 2, 3, 1)).astype('uint8') - if args.times == 4: - mid0 = ((mid0[:, :, :h, :w] * 255.).cpu().detach().numpy().transpose(0, 2, 3, 1)).astype('uint8') - mid2 = ((mid2[:, :, :h, :w] * 255.).cpu().detach().numpy().transpose(0, 2, 3, 1)).astype('uint8') - else: - mid0, mid2 = None, None - writeframe(I0, mid0, mid1, mid2, I1, p.mean(3).mean(2).mean(1)) + inferences = list(map(lambda x: ((x[:, :, :h, :w] * 255.).cpu().detach().numpy().transpose(0, 2, 3, 1)).astype('uint8'), inferences)) + + write_frame(vid_out, I0, inferences, I1, p.mean(3).mean(2).mean(1), args) pbar.update(4) img_list = img_list[-1:] pbar.close() +vid_out.release() \ No newline at end of file