From d57665025fe8f7bb84e94dfdcc0fc0d06a24bbf9 Mon Sep 17 00:00:00 2001
From: hzwer <598460606@163.com>
Date: Mon, 23 Nov 2020 11:30:46 +0800
Subject: [PATCH] Count frame

---
 inference_video_parallel.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/inference_video_parallel.py b/inference_video_parallel.py
index 1c397bc..9bca3cc 100644
--- a/inference_video_parallel.py
+++ b/inference_video_parallel.py
@@ -108,9 +108,11 @@ print('{}.{}, {} frames in total, {}FPS to {}FPS'.format(video_path_wo_ext, args
 pbar = tqdm(total=tot_frame)
 img_list = []
 _thread.start_new_thread(clear_buffer, (args, ))
+tot_frame -= 1
 for frame in videogen:
+    tot_frame -= 1
     img_list.append(frame)
-    if len(img_list) == 5 or (not success and len(img_list) > 1):
+    if len(img_list) == 5 or tot_frame == 0:
         imgs = torch.from_numpy(np.transpose(img_list, (0, 3, 1, 2))).to(device, non_blocking=True).float() / 255.
         I0 = imgs[:-1]
         I1 = imgs[1:]