Mirrors/ECCV2022-RIFE
1
0
Fork 0
mirror of https://github.com/hzwer/ECCV2022-RIFE.git synced 2026-02-24 04:19:41 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add montage

Browse Source
This commit is contained in:
hzwer
2020-11-15 16:41:52 +08:00
parent 0a0c10052b
commit a73928eebc
1 changed files with 17 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
inference_mp4_2x.py
View File

@@ -10,6 +10,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
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')
args = parser.parse_args()
model = Model()
@@ -26,7 +27,10 @@ pw = ((w - 1) // 32 + 1) * 32
padding = (0, pw - w, 0, ph - h)
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
print('{}.mp4'.format(args.video[:-4]))
output = cv2.VideoWriter('{}_2x.mp4'.format(args.video[:-4]), fourcc, fps * 2, (w, h))
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
while success:
lastframe = frame
@@ -36,13 +40,20 @@ while success:
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', recompute_scale_factor=False)
- F.interpolate(I1, (16, 16), mode='bilinear', recompute_scale_factor=False)).abs().mean() > 0.2:
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
else:
mid = model.inference(I0, I1)
mid = ((mid[0].cpu().detach().numpy().transpose(1, 2, 0))*255.).astype('uint8')
output.write(lastframe[:h, :w])
output.write(mid[:h, :w])
output.write(frame)
if args.montage:
output.write(torch.cat((lastframe, lastframe), 2))
output.write(torch.cat((lastframe, mid[:h, :w]), 2))
else:
output.write(lastframe)
output.write(mid[:h, :w])
if args.montage:
output.write(torch.cat((frame, frame), 2))
else:
output.write(frame)
output.release()