From 9df30cdcf3395941d343d9be13a1628b0040b1c9 Mon Sep 17 00:00:00 2001
From: hzwer <598460606@163.com>
Date: Sun, 15 Nov 2020 17:33:44 +0800
Subject: [PATCH] Fix 60fps

---
 inference_mp4_2x.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/inference_mp4_2x.py b/inference_mp4_2x.py
index 804f06e..6acea8f 100644
--- a/inference_mp4_2x.py
+++ b/inference_mp4_2x.py
@@ -16,6 +16,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)
 args = parser.parse_args()
 
 model = Model()
@@ -24,7 +25,7 @@ model.eval()
 model.device()
 
 videoCapture = cv2.VideoCapture(args.video)
-fps = videoCapture.get(cv2.CAP_PROP_FPS)
+fps = np.round(videoCapture.get(cv2.CAP_PROP_FPS))
 success, frame = videoCapture.read()
 h, w, _ = frame.shape
 ph = ((h - 1) // 32 + 1) * 32
@@ -32,12 +33,12 @@ 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, fps * 2))
+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))
+    output = cv2.VideoWriter('{}_2x.mp4'.format(args.video[:-4]), fourcc, args.fps, (2*w, h))
 else:        
-    output = cv2.VideoWriter('{}_2x.mp4'.format(args.video[:-4]), fourcc, fps * 2, (w, h))
+    output = cv2.VideoWriter('{}_2x.mp4'.format(args.video[:-4]), fourcc, args.fps, (w, h))
 frame = frame
 while success:
     lastframe = frame