diff --git a/inference_video_parallel.py b/inference_video_parallel.py index 5d035f2..de904ab 100644 --- a/inference_video_parallel.py +++ b/inference_video_parallel.py @@ -104,7 +104,7 @@ while success: success, frame = videoCapture.read() if success: img_list.append(frame) - if len(img_list) == 3 or (not success and len(img_list) > 1): + if len(img_list) == 5 or (not success and len(img_list) > 1): I0 = torch.from_numpy(np.transpose(img_list[:-1], (0, 3, 1, 2)).astype('float32') / 255.).to(device, non_blocking=True) I1 = torch.from_numpy(np.transpose(img_list[1:], (0, 3, 1, 2)).astype('float32') / 255.).to(device, non_blocking=True) p = (F.interpolate(I0, (16, 16), mode='bilinear', align_corners=False)