diff --git a/inference_mp4_2x.py b/inference_mp4_2x.py index e47d6ea..01a0b68 100644 --- a/inference_mp4_2x.py +++ b/inference_mp4_2x.py @@ -14,8 +14,9 @@ if torch.cuda.is_available(): parser = argparse.ArgumentParser(description='Interpolation for a pair of images') parser.add_argument('--video', dest='video', required=True) -parser.add_argument('--model', dest='model', type=str, default='RIFE') parser.add_argument('--montage', dest='montage', action='store_true', help='montage origin video') +parser.add_argument('--fps', dest='fps', type=int, default=60) +parser.add_argument('--model', dest='model', type=str, default='RIFE') args = parser.parse_args() if args.model == '2F': @@ -31,38 +32,41 @@ videoCapture = cv2.VideoCapture(args.video) fps = np.round(videoCapture.get(cv2.CAP_PROP_FPS)) success, frame = videoCapture.read() h, w, _ = frame.shape +fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') +output = cv2.VideoWriter('{}_4x.mp4'.format(args.video[:-4]), fourcc, args.fps, (w, h)) +if args.montage: + left = w // 4 + w = w // 2 ph = ((h - 1) // 32 + 1) * 32 pw = ((w - 1) // 32 + 1) * 32 padding = (0, pw - w, 0, ph - h) -fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT) -print('{}.mp4, {} frames in total, {}FPS to {}FPS'.format(args.video[:-4], tot_frame, fps, 2*fps)) +print('{}.mp4, {} frames in total, {}FPS to {}FPS'.format(args.video[:-4], tot_frame, fps, args.fps)) pbar = tqdm(total=tot_frame) if args.montage: - output = cv2.VideoWriter('{}_2x.mp4'.format(args.video[:-4]), fourcc, fps*2, (2*w, h)) -else: - output = cv2.VideoWriter('{}_2x.mp4'.format(args.video[:-4]), fourcc, fps*2, (w, h)) -frame = frame + frame = frame[:, left: left + w] while success: lastframe = frame success, frame = videoCapture.read() + if args.montage: + frame = frame[:, left: left + w] if success: I0 = torch.from_numpy(np.transpose(lastframe, (2,0,1)).astype("float32") / 255.).to(device).unsqueeze(0) I1 = torch.from_numpy(np.transpose(frame, (2,0,1)).astype("float32") / 255.).to(device).unsqueeze(0) I0 = F.pad(I0, padding) I1 = F.pad(I1, padding) if (F.interpolate(I0, (16, 16), mode='bilinear', align_corners=False) - - F.interpolate(I1, (16, 16), mode='bilinear', align_corners=False)).abs().mean() > 0.2: - mid = lastframe + - F.interpolate(I1, (16, 16), mode='bilinear', align_corners=False)).abs().mean() > 0.2: + mid1 = lastframe else: - mid = model.inference(I0, I1) - mid = (((mid[0] * 255.).cpu().detach().numpy().transpose(1, 2, 0))).astype('uint8') + mid1 = model.inference(I0, I1) + mid1 = (((mid1[0] * 255.).cpu().detach().numpy().transpose(1, 2, 0))).astype('uint8') if args.montage: output.write(np.concatenate((lastframe, lastframe), 1)) - output.write(np.concatenate((lastframe, mid[:h, :w]), 1)) + output.write(np.concatenate((lastframe, mid1[:h, :w]), 1)) else: output.write(lastframe) - output.write(mid[:h, :w]) + output.write(mid1[:h, :w]) pbar.update(1) if args.montage: output.write(np.concatenate((lastframe, lastframe), 1)) diff --git a/inference_mp4_4x.py b/inference_mp4_4x.py index 80d3b3a..f1b2624 100644 --- a/inference_mp4_4x.py +++ b/inference_mp4_4x.py @@ -15,6 +15,7 @@ if torch.cuda.is_available(): parser = argparse.ArgumentParser(description='Interpolation for a pair of images') parser.add_argument('--video', dest='video', required=True) parser.add_argument('--montage', dest='montage', action='store_true', help='montage origin video') +parser.add_argument('--fps', dest='fps', type=int, default=60) parser.add_argument('--model', dest='model', type=str, default='RIFE') args = parser.parse_args() @@ -31,21 +32,24 @@ videoCapture = cv2.VideoCapture(args.video) fps = np.round(videoCapture.get(cv2.CAP_PROP_FPS)) success, frame = videoCapture.read() h, w, _ = frame.shape +fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') +output = cv2.VideoWriter('{}_4x.mp4'.format(args.video[:-4]), fourcc, args.fps, (w, h)) +if args.montage: + left = w // 4 + w = w // 2 ph = ((h - 1) // 32 + 1) * 32 pw = ((w - 1) // 32 + 1) * 32 padding = (0, pw - w, 0, ph - h) -fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT) -print('{}.mp4, {} frames in total, {}FPS to {}FPS'.format(args.video[:-4], tot_frame, fps, 4*fps)) +print('{}.mp4, {} frames in total, {}FPS to {}FPS'.format(args.video[:-4], tot_frame, fps, args.fps)) pbar = tqdm(total=tot_frame) if args.montage: - output = cv2.VideoWriter('{}_4x.mp4'.format(args.video[:-4]), fourcc, fps*4, (2*w, h)) -else: - output = cv2.VideoWriter('{}_4x.mp4'.format(args.video[:-4]), fourcc, fps*4, (w, h)) -frame = frame + frame = frame[:, left: left + w] while success: lastframe = frame success, frame = videoCapture.read() + if args.montage: + frame = frame[:, left: left + w] if success: I0 = torch.from_numpy(np.transpose(lastframe, (2,0,1)).astype("float32") / 255.).to(device).unsqueeze(0) I1 = torch.from_numpy(np.transpose(frame, (2,0,1)).astype("float32") / 255.).to(device).unsqueeze(0) @@ -58,11 +62,10 @@ while success: mid2 = frame else: mid1 = model.inference(I0, I1) - mid0 = model.inference(I0, mid1) - mid2 = model.inference(mid1, I1) - mid0 = (((mid0[0] * 255.).cpu().detach().numpy().transpose(1, 2, 0))).astype('uint8') + mid = model.inference(torch.cat((I0, mid1), 0), torch.cat((mid1, I1), 0)) + mid0 = (((mid[0] * 255.).cpu().detach().numpy().transpose(1, 2, 0))).astype('uint8') mid1 = (((mid1[0] * 255.).cpu().detach().numpy().transpose(1, 2, 0))).astype('uint8') - mid2 = (((mid2[0] * 255.).cpu().detach().numpy().transpose(1, 2, 0))).astype('uint8') + mid2 = (((mid[1]* 255.).cpu().detach().numpy().transpose(1, 2, 0))).astype('uint8') if args.montage: output.write(np.concatenate((lastframe, lastframe), 1)) output.write(np.concatenate((lastframe, mid0[:h, :w]), 1))