Convert video to RGB on the CPU with the matrix of the stream

DALI_EXTRA_VERSION

@@ -1 +1 @@
-7fa5101b6ffaf536e2e9edb34bf5bea7ed0a74ad
+eda69271b798e84648e3c51445ce9cf183c36c1e

dali/operators/video/frames_decoder_cpu.cc

@@ -107,10 +107,19 @@ void FramesDecoderCpu::CopyToOutput(uint8_t *data) {
   bool src_full_range = frame_->color_range == AVCOL_RANGE_JPEG ||
                         (frame_->color_range == AVCOL_RANGE_UNSPECIFIED &&
                          codec_params_->color_range == AVCOL_RANGE_JPEG);
-  if (sws_src_full_range_ != src_full_range) {
+  // The matrix that the stream says it uses (BT.709 for most HD video); otherwise, BT.601.
+  // It matters only when converting to RGB: YCbCr output is a copy of the decoded planes, which
+  // must not be converted to another matrix (the destination one is BT.601).
+  AVColorSpace src_colorspace = frame_->colorspace != AVCOL_SPC_UNSPECIFIED
+                                    ? frame_->colorspace
+                                    : codec_params_->color_space;
+  int sws_colorspace = (image_type_ == DALI_RGB && src_colorspace == AVCOL_SPC_BT709)
+                           ? SWS_CS_ITU709
+                           : SWS_CS_DEFAULT;
+  if (sws_src_full_range_ != src_full_range || sws_src_colorspace_ != sws_colorspace) {
     int ret = sws_setColorspaceDetails(
       sws_ctx_.get(),
-      sws_getCoefficients(SWS_CS_DEFAULT),
+      sws_getCoefficients(sws_colorspace),
       src_full_range,
       sws_getCoefficients(SWS_CS_DEFAULT),
       1,
@@ -121,6 +130,7 @@ void FramesDecoderCpu::CopyToOutput(uint8_t *data) {
                  make_string("Could not set color space conversion details: ",
                              av_error_string(ret)));
     sws_src_full_range_ = src_full_range;
+    sws_src_colorspace_ = sws_colorspace;
   }
 
   uint8_t *dest[4] = {sws_output_data, nullptr, nullptr, nullptr};

dali/operators/video/frames_decoder_cpu.h

@@ -65,6 +65,7 @@ class DLL_PUBLIC FramesDecoderCpu : public FramesDecoderBase {
   std::unique_ptr<SwsContext, decltype(&sws_freeContext)> sws_ctx_{
     nullptr, sws_freeContext};
   std::optional<bool> sws_src_full_range_;
+  std::optional<int> sws_src_colorspace_;
 
   std::vector<uint8_t> tmp_buffer_;
 };

dali/test/python/decoder/test_video.py

@@ -306,6 +306,63 @@ def test_pipeline():
     assert np.mean(absdiff) < 2
 
 
+COLOR_MATRIX_NUM_BARS = 10
+COLOR_MATRIX_BAR_WIDTH = 192  # pixels
+
+
+def color_matrix_bar(frame, i):
+    """Mean color of the interior of the i-th bar of the test pattern in db/video/color_matrix: 10
+    flat bars (red, green, blue, yellow, cyan, magenta, orange, grey, black, white), 192 pixels
+    wide"""
+    interior = (
+        slice(100, 440),
+        slice(i * COLOR_MATRIX_BAR_WIDTH + 32, (i + 1) * COLOR_MATRIX_BAR_WIDTH - 32),
+    )
+    return frame[interior].reshape(-1, 3).mean(axis=0)
+
+
+def color_matrix_bars_error(frame, reference):
+    """Worst difference between `frame` and `reference` of the mean colors of the interiors of the
+    bars of the test pattern in db/video/color_matrix"""
+    return max(
+        np.max(np.abs(color_matrix_bar(frame, i) - color_matrix_bar(reference, i)))
+        for i in range(COLOR_MATRIX_NUM_BARS)
+    )
+
+
+def check_color_matrix_black_and_white(frame):
+    """Black and white are the ends of the range (Y=16 and Y=235 in the limited range): the luma
+    has to be mapped to the whole range and clamped there (not 16 and 235). They do not depend on
+    the matrix, unlike the saturated bars."""
+    black, white = color_matrix_bar(frame, 8), color_matrix_bar(frame, 9)
+    assert np.all(black <= 2), f"black bar is {black}"
+    assert np.all(white >= 253), f"white bar is {white}"
+
+
+@params("bt709_limited_1080p_vp9")
+def test_color_matrix_bt709_cpu(name):
+    # The streams are tagged as BT.709 (limited range). The references are what ffmpeg decodes,
+    # which uses the matrix from the stream; converting with BT.601 gives errors up to 22 in the
+    # interior of the saturated bars.
+    data_dir = os.path.join(get_dali_extra_path(), "db", "video", "color_matrix")
+
+    @pipeline_def
+    def test_pipeline():
+        return fn.experimental.readers.video(
+            device="cpu",
+            filenames=[os.path.join(data_dir, f"{name}.mp4")],
+            sequence_length=1,
+        )
+
+    out = test_pipeline(batch_size=1, num_threads=1, device_id=None).run()[0]
+    frame = out.as_array()[0, 0]
+    ref = cv2.cvtColor(cv2.imread(os.path.join(data_dir, f"{name}_ref.png")), cv2.COLOR_BGR2RGB)
+    assert frame.shape == ref.shape
+    assert color_matrix_bars_error(frame, ref) <= 4
+    check_color_matrix_black_and_white(frame)
+    assert np.mean(np.abs(frame.astype(int) - ref.astype(int))) < 2
+
+
 @attr("multi_gpu")
 @params("cpu", "mixed")
 def test_multi_gpu_video(device):