From c925f9e5ee0506849ad5b5b9980b6a5e3fc3248d Mon Sep 17 00:00:00 2001 From: QuentinDuval Date: Fri, 19 Jun 2020 05:23:38 -0700 Subject: [PATCH] Fix: PyAV does not support floating point numbers with decimals as FPS when writing and will throw in case this constraint is not satisfied. --- torchvision/io/video.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/torchvision/io/video.py b/torchvision/io/video.py index f822cbd3343..2959024b5dd 100644 --- a/torchvision/io/video.py +++ b/torchvision/io/video.py @@ -2,7 +2,7 @@ import math import re import warnings -from typing import Tuple, List +from typing import List, Tuple, Union import numpy as np import torch @@ -49,7 +49,7 @@ def _av_available(): _GC_COLLECTION_INTERVAL = 10 -def write_video(filename, video_array, fps, video_codec="libx264", options=None): +def write_video(filename, video_array, fps: Union[int, float], video_codec="libx264", options=None): """ Writes a 4d tensor in [T, H, W, C] format in a video file @@ -65,6 +65,11 @@ def write_video(filename, video_array, fps, video_codec="libx264", options=None) _check_av_available() video_array = torch.as_tensor(video_array, dtype=torch.uint8).numpy() + # PyAV does not support floating point numbers with decimal point + # and will throw OverflowException in case this is not the case + if isinstance(fps, float): + fps = np.round(fps) + container = av.open(filename, mode="w") stream = container.add_stream(video_codec, rate=fps)