diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ce37ff4f9..2a8d815b7 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -39,9 +39,12 @@ Features: - Support reusing the thread's current CUDA context via a ``current_ctx`` flag on ``CudaContext`` and ``VideoFrame.from_dlpack``, for interop with libraries like PyTorch that initialize CUDA first by :gh-user:`Yozer` (:pr:`2339`). - ``VideoFrame.from_dlpack`` no longer requires restating ``primary_ctx``/``current_ctx`` when passing an explicit ``cuda_context``; the flags are only validated when explicitly given by :gh-user:`WyattBlue`. - Support passing an explicit CUDA stream to FFmpeg CUDA operations, including NVENC input and output, via a ``cuda_stream`` parameter on ``CudaContext``; currently limited to logical CUDA device 0 by :gh-user:`Yozer` (:pr:`2360`). +- ``VideoFrame.save`` now forwards keyword arguments to the encoder, letting callers trade file size for speed (e.g. ``pred="none"`` or ``compression_level=1`` for PNG, ``qscale=2`` for JPG) by :gh-user:`WyattBlue`. Fixes: +- Fix ``VideoFrame.save`` raising ``AttributeError`` when given a ``Path`` rather than a ``str`` by :gh-user:`WyattBlue`. + - Prevent crashes and corrupted output when structural codec properties are changed after an output stream has been opened by :gh-user:`WyattBlue`, reported by :gh-user:`oakaigh` (:issue:`2232`). - Fix a crash when using a stream that has no ``CodecContext`` (a demuxed stream with no available decoder, such as one from a truncated file, or a stream created by ``add_mux_stream``); decoding now raises ``DecoderNotFoundError``, encoding now raises ``EncoderNotFoundError``, and ``BitStreamFilterContext`` accepts such a stream as ``out_stream`` by :gh-user:`WyattBlue`, reported by :gh-user:`justinrmiller` (:issue:`2344`). diff --git a/av/video/frame.pxd b/av/video/frame.pxd index 66926058d..7f93c99ec 100644 --- a/av/video/frame.pxd +++ b/av/video/frame.pxd @@ -25,6 +25,5 @@ cdef class VideoFrame(Frame): cdef readonly int _device_id cdef _init(self, lib.AVPixelFormat format, unsigned int width, unsigned int height) cdef _init_user_attributes(self) - cpdef save(self, object filepath) cdef VideoFrame alloc_video_frame() diff --git a/av/video/frame.py b/av/video/frame.py index 74964107b..0d22c481c 100644 --- a/av/video/frame.py +++ b/av/video/frame.py @@ -773,17 +773,21 @@ def to_rgb(self, **kwargs): """ return self.reformat(format="rgb24", **kwargs) - @cython.ccall - def save(self, filepath: object): + def save(self, filepath: object, **options): """Save a VideoFrame as a JPG or PNG. :param filepath: str | Path + :param \\**options: Encoder options, e.g. ``pred="none"`` or + ``compression_level=1`` for PNG, ``qscale=2`` for JPG. Values are + coerced to ``str``. The PNG defaults favor file size over speed; + ``pred="none"`` is roughly 3x faster and 2x larger. """ is_jpg: cython.bint + name: str = str(filepath) - if filepath.endswith(".png"): + if name.endswith(".png"): is_jpg = False - elif filepath.endswith(".jpg") or filepath.endswith(".jpeg"): + elif name.endswith(".jpg") or name.endswith(".jpeg"): is_jpg = True else: raise ValueError("filepath must end with png or jpg.") @@ -794,7 +798,11 @@ def save(self, filepath: object): from av.container.core import open with open(filepath, "w", options={"update": "1"}) as output: - output_stream = output.add_stream(encoder, pix_fmt=pix_fmt) + output_stream = output.add_stream( + encoder, + pix_fmt=pix_fmt, + options={k: str(v) for k, v in options.items()}, + ) output_stream.width = self.width output_stream.height = self.height diff --git a/av/video/frame.pyi b/av/video/frame.pyi index f07a8079f..b2240ab9c 100644 --- a/av/video/frame.pyi +++ b/av/video/frame.pyi @@ -85,7 +85,7 @@ class VideoFrame(Frame): threads: int | None = None, ) -> VideoFrame: ... def to_rgb(self, **kwargs: Any) -> VideoFrame: ... - def save(self, filepath: str | Path) -> None: ... + def save(self, filepath: str | Path, **options: Any) -> None: ... def to_image(self, **kwargs): ... def to_ndarray( self, channel_last: bool = False, **kwargs: Any diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index acd3f53ab..36a1d5386 100644 --- a/tests/test_videoframe.py +++ b/tests/test_videoframe.py @@ -1379,3 +1379,21 @@ def test_reformat_pixel_format_align() -> None: result = frame_rgb.to_ndarray() assert result.shape == expected_rgb.shape assert numpy.abs(result.astype(int) - expected_rgb.astype(int)).max() <= 1 + + +def test_save_options(tmp_path) -> None: + y, x = numpy.mgrid[0:240, 0:320] + array = numpy.dstack([x % 256, y % 256, (x + y) % 256]).astype(numpy.uint8) + frame = VideoFrame.from_ndarray(array, format="rgb24") + + default = tmp_path / "default.png" + unfiltered = tmp_path / "unfiltered.png" + frame.save(default) + frame.save(unfiltered, pred="none") + + # pred="none" skips the row filter: much faster, much bigger. + assert unfiltered.stat().st_size > 4 * default.stat().st_size + + # Non-str values are coerced. + frame.save(tmp_path / "q.jpg", qscale=2) + assert (tmp_path / "q.jpg").stat().st_size > 0