Add binding for "avcodec_find_best_pix_fmt_of_list" (#2058)

Author: ndeybach

av/codec/__init__.py

@@ -1,4 +1,11 @@
-from .codec import Capabilities, Codec, Properties, codec_descriptor, codecs_available
+from .codec import (
+    Capabilities,
+    Codec,
+    Properties,
+    codec_descriptor,
+    codecs_available,
+    find_best_pix_fmt_of_list,
+)
 from .context import CodecContext
 
 __all__ = (
@@ -7,5 +14,6 @@
     "Properties",
     "codec_descriptor",
     "codecs_available",
+    "find_best_pix_fmt_of_list",
     "CodecContext",
 )

av/codec/codec.pyi

@@ -1,6 +1,6 @@
 from enum import Flag, IntEnum
 from fractions import Fraction
-from typing import ClassVar, Literal, cast, overload
+from typing import ClassVar, Literal, Sequence, cast, overload
 
 from av.audio.codeccontext import AudioCodecContext
 from av.audio.format import AudioFormat
@@ -113,3 +113,27 @@ codecs_available: set[str]
 
 def dump_codecs() -> None: ...
 def dump_hwconfigs() -> None: ...
+
+PixFmtLike = str | VideoFormat
+
+def find_best_pix_fmt_of_list(
+    pix_fmts: Sequence[PixFmtLike],
+    src_pix_fmt: PixFmtLike,
+    has_alpha: bool = False,
+) -> tuple[VideoFormat | None, int]:
+    """
+    Find the best pixel format to convert to given a source format.
+
+    Wraps :ffmpeg:`avcodec_find_best_pix_fmt_of_list`.
+
+    :param pix_fmts: Iterable of pixel formats to choose from (str or VideoFormat).
+    :param src_pix_fmt: Source pixel format (str or VideoFormat).
+    :param bool has_alpha: Whether the source alpha channel is used.
+    :return: (best_format, loss): best_format is the best matching pixel format from
+        the list, or None if no suitable format was found; loss is Combination of flags informing you what kind of losses will occur.
+    :rtype: (VideoFormat | None, int)
+
+    Note on loss: it is a bitmask of FFmpeg loss flags describing what kinds of information would be lost converting from src_pix_fmt to best_format (e.g. loss of alpha, chroma, colorspace, resolution, bit depth, etc.). Multiple losses can be present at once, so the value is meant to be interpreted with bitwise & against FFmpeg's FF_LOSS_* constants.
+    For exact behavior see: libavutil/pixdesc.c/get_pix_fmt_score() in ffmpeg source code.
+    """
+    ...

av/codec/codec.pyx

@@ -4,9 +4,10 @@ from av.audio.format cimport get_audio_format
 from av.codec.hwaccel cimport wrap_hwconfig
 from av.descriptor cimport wrap_avclass
 from av.utils cimport avrational_to_fraction
-from av.video.format cimport get_video_format
+from av.video.format cimport VideoFormat, get_pix_fmt, get_video_format
 
 from enum import Flag, IntEnum
+from libc.stdlib cimport free, malloc
 
 
 cdef object _cinit_sentinel = object()
@@ -387,3 +388,59 @@ def dump_hwconfigs():
         print("   ", codec.name)
         for config in configs:
             print("       ", config)
+
+
+def find_best_pix_fmt_of_list(pix_fmts, src_pix_fmt, has_alpha=False):
+    """
+    Find the best pixel format to convert to given a source format.
+
+    Wraps :ffmpeg:`avcodec_find_best_pix_fmt_of_list`.
+
+    :param pix_fmts: Iterable of pixel formats to choose from (str or VideoFormat).
+    :param src_pix_fmt: Source pixel format (str or VideoFormat).
+    :param bool has_alpha: Whether the source alpha channel is used.
+    :return: (best_format, loss)
+    :rtype: (VideoFormat | None, int)
+    """
+    cdef lib.AVPixelFormat src
+    cdef lib.AVPixelFormat best
+    cdef lib.AVPixelFormat *c_list = NULL
+    cdef Py_ssize_t n
+    cdef Py_ssize_t i
+    cdef object item
+    cdef int c_loss
+
+    if pix_fmts is None:
+        raise TypeError("pix_fmts must not be None")
+
+    pix_fmts = tuple(pix_fmts)
+    if not pix_fmts:
+        return None, 0
+
+    if isinstance(src_pix_fmt, VideoFormat):
+        src = (<VideoFormat>src_pix_fmt).pix_fmt
+    else:
+        src = get_pix_fmt(<str>src_pix_fmt)
+
+    n = len(pix_fmts)
+    c_list = <lib.AVPixelFormat *>malloc((n + 1) * sizeof(lib.AVPixelFormat))
+    if c_list == NULL:
+        raise MemoryError()
+
+    try:
+        for i in range(n):
+            item = pix_fmts[i]
+            if isinstance(item, VideoFormat):
+                c_list[i] = (<VideoFormat>item).pix_fmt
+            else:
+                c_list[i] = get_pix_fmt(<str>item)
+        c_list[n] = lib.AV_PIX_FMT_NONE
+
+        c_loss = 0
+        best = lib.avcodec_find_best_pix_fmt_of_list(
+            c_list, src, 1 if has_alpha else 0, &c_loss
+        )
+        return get_video_format(best, 0, 0), c_loss
+    finally:
+        if c_list != NULL:
+            free(c_list)

include/libavcodec/avcodec.pxd

@@ -75,6 +75,13 @@ cdef extern from "libavcodec/avcodec.h" nogil:
     cdef char* avcodec_configuration()
     cdef char* avcodec_license()
 
+    AVPixelFormat avcodec_find_best_pix_fmt_of_list(
+        const AVPixelFormat *pix_fmt_list,
+        AVPixelFormat src_pix_fmt,
+        int has_alpha,
+        int *loss_ptr,
+    )
+
     cdef size_t AV_INPUT_BUFFER_PADDING_SIZE
     cdef int64_t AV_NOPTS_VALUE
 

tests/test_codec.py

@@ -1,6 +1,7 @@
 import pytest
 
 from av import AudioFormat, Codec, VideoFormat, codecs_available
+from av.codec import find_best_pix_fmt_of_list
 from av.codec.codec import UnknownCodecError
 
 
@@ -89,3 +90,64 @@ def test_codec_opus_encoder() -> None:
 
 def test_codecs_available() -> None:
     assert codecs_available
+
+
+def test_find_best_pix_fmt_of_list_empty() -> None:
+    best, loss = find_best_pix_fmt_of_list([], "rgb24")
+    assert best is None
+    assert loss == 0
+
+
+@pytest.mark.parametrize(
+    "pix_fmts,src_pix_fmt,expected_best",
+    [
+        (["rgb24", "yuv420p"], "rgb24", "rgb24"),
+        (["rgb24"], "yuv420p", "rgb24"),
+        (["yuv420p"], "rgb24", "yuv420p"),
+        ([VideoFormat("yuv420p")], VideoFormat("rgb24"), "yuv420p"),
+        (
+            ["yuv420p", "yuv444p", "gray", "rgb24", "rgba", "bgra", "yuyv422"],
+            "rgba",
+            "rgba",
+        ),
+    ],
+)
+def test_find_best_pix_fmt_of_list_best(pix_fmts, src_pix_fmt, expected_best) -> None:
+    best, loss = find_best_pix_fmt_of_list(pix_fmts, src_pix_fmt)
+    assert best is not None
+    assert best.name == expected_best
+    assert isinstance(loss, int)
+
+
+@pytest.mark.parametrize(
+    "pix_fmts,src_pix_fmt",
+    [
+        (["__unknown_pix_fmt"], "rgb24"),
+        (["rgb24"], "__unknown_pix_fmt"),
+    ],
+)
+def test_find_best_pix_fmt_of_list_unknown_pix_fmt(pix_fmts, src_pix_fmt) -> None:
+    with pytest.raises(ValueError, match="not a pixel format"):
+        find_best_pix_fmt_of_list(pix_fmts, src_pix_fmt)
+
+
+@pytest.mark.parametrize(
+    "pix_fmts,src_pix_fmt",
+    [
+        (["rgb24", "bgr24", "gray", "yuv420p", "yuv444p", "yuyv422"], "nv12"),
+        (["yuv420p", "yuv444p", "gray", "yuv420p"], "rgb24"),
+        (["rgb24", "rgba", "bgra", "rgb24", "gray"], "yuv420p"),
+    ],
+)
+def test_find_best_pix_fmt_of_list_picks_from_list(pix_fmts, src_pix_fmt) -> None:
+    best, loss = find_best_pix_fmt_of_list(pix_fmts, src_pix_fmt)
+    assert best is not None
+    assert best.name in set(pix_fmts)
+    assert isinstance(loss, int)
+
+
+def test_find_best_pix_fmt_of_list_alpha_loss_flagged_when_used() -> None:
+    best, loss = find_best_pix_fmt_of_list(["rgb24"], "rgba", has_alpha=True)
+    assert best is not None
+    assert best.name == "rgb24"
+    assert loss != 0