Allow configuring hardware encoder sw_format

Author: CarlosRDomin

av/codec/context.py

@@ -396,9 +396,16 @@ def _setup_encode_hwframes(self) -> cython.void:
             return  # Already set up.
 
         hw_format: lib.AVPixelFormat = self.hwaccel_ctx.config.ptr.pix_fmt
-        sw_format: lib.AVPixelFormat = cython.cast(lib.AVPixelFormat, self.ptr.pix_fmt)
+        sw_format: lib.AVPixelFormat = cython.cast(
+            lib.AVPixelFormat, self.ptr.sw_pix_fmt
+        )
+
+        # The codec context's sw_pix_fmt holds the software format the user
+        # wants the hardware frames context to use. Fall back to pix_fmt to
+        # preserve the existing stream.pix_fmt configuration path.
+        if sw_format == lib.AV_PIX_FMT_NONE:
+            sw_format = cython.cast(lib.AVPixelFormat, self.ptr.pix_fmt)
 
-        # The codec context's pix_fmt holds the *software* format the user feeds in.
         # If they left it as the hardware format (or unset), pick a sane default.
         if sw_format == hw_format or sw_format == lib.AV_PIX_FMT_NONE:
             sw_format = lib.av_get_pix_fmt(b"nv12")

av/video/codeccontext.py

@@ -261,12 +261,22 @@ def sw_format(self):
         :type: VideoFormat | None
         """
         if not self.ptr.hw_frames_ctx:
+            if self.ptr.sw_pix_fmt != lib.AV_PIX_FMT_NONE:
+                return get_video_format(
+                    cython.cast(lib.AVPixelFormat, self.ptr.sw_pix_fmt),
+                    self.ptr.width,
+                    self.ptr.height,
+                )
             return None
         frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
             cython.pointer[lib.AVHWFramesContext], self.ptr.hw_frames_ctx.data
         )
         return get_video_format(frames_ctx.sw_format, self.ptr.width, self.ptr.height)
 
+    @sw_format.setter
+    def sw_format(self, value):
+        self.ptr.sw_pix_fmt = get_pix_fmt(value)
+
     @property
     def framerate(self):
         """

tests/test_encode.py

@@ -517,6 +517,13 @@ def test_profiles(self) -> None:
 }
 
 
+def get_hwaccel_format(encoder: str, device_type: str) -> str:
+    for config in av.Codec(encoder, "w").hardware_configs:
+        if config.device_type.name == device_type and config.format is not None:
+            return config.format.name
+    pytest.skip(f"No hardware format for {device_type} on {encoder}")
+
+
 def test_hardware_encode() -> None:
     hwdevices_available = av.codec.hwaccel.hwdevices_available()
     if "HWACCEL_DEVICE_TYPE" not in os.environ:
@@ -569,3 +576,42 @@ def test_hardware_encode() -> None:
     with av.open(file, "r") as in_container:
         decoded = sum(1 for _ in in_container.decode(video=0))
     assert decoded == n_frames
+
+
+def test_hardware_encode_honors_sw_format() -> None:
+    hwdevices_available = av.codec.hwaccel.hwdevices_available()
+    if "HWACCEL_DEVICE_TYPE" not in os.environ:
+        pytest.skip(
+            "Set the HWACCEL_DEVICE_TYPE to run this test. "
+            f"Options are {' '.join(hwdevices_available)}"
+        )
+
+    device_type = os.environ["HWACCEL_DEVICE_TYPE"]
+    assert device_type in hwdevices_available, f"{device_type} not available"
+
+    encoder = _HWACCEL_ENCODERS.get(device_type)
+    if encoder is None:
+        pytest.skip(f"No hardware encoder mapped for {device_type}")
+    hw_format = get_hwaccel_format(encoder, device_type)
+
+    hwaccel = av.codec.hwaccel.HWAccel(
+        device_type=device_type, allow_software_fallback=False
+    )
+    container = av.open(io.BytesIO(), mode="w", format="mp4")
+    stream = container.add_stream(encoder, rate=30, hwaccel=hwaccel)
+    stream.width = 320
+    stream.height = 240
+    stream.pix_fmt = hw_format
+    stream.codec_context.sw_format = "yuv420p"
+
+    assert stream.codec_context.sw_format.name == "yuv420p"
+
+    frame = VideoFrame(320, 240, "rgb24")
+    for packet in stream.encode(frame):
+        container.mux(packet)
+
+    assert stream.codec_context.pix_fmt == hw_format
+    assert stream.codec_context.sw_format.name == "yuv420p"
+    for packet in stream.encode():
+        container.mux(packet)
+    container.close()