comments

Author: Dan-Flores

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -373,7 +373,7 @@ std::string CudaDeviceInterface::getDetails() {
 // Below are methods exclusive to video encoding:
 // --------------------------------------------------------------------------
 namespace {
-// RGB to YUV conversion matrices organized by color range and color space
+// RGB to YUV conversion matrices to use in NPP color conversion functions
 struct ColorConversionMatrices {
   static constexpr Npp32f BT601_LIMITED[3][4] = {
       {0.257f, 0.504f, 0.098f, 16.0f},
@@ -406,19 +406,17 @@ struct ColorConversionMatrices {
       {0.5f, -0.459786f, -0.040214f, 128.0f}};
 };
 
-// Returns the appropriate conversion matrix based on codec context color space
-// and range
+// Returns conversion matrix based on codec context color space and range
 const Npp32f (*getConversionMatrix(AVCodecContext* codecContext))[4] {
-  // Limited range (TV/Video range 16-235) - default if unspecified
-  if (codecContext->color_range == AVCOL_RANGE_MPEG ||
+  if (codecContext->color_range == AVCOL_RANGE_MPEG || // limited range
       codecContext->color_range == AVCOL_RANGE_UNSPECIFIED) {
     if (codecContext->colorspace == AVCOL_SPC_BT470BG) {
       return ColorConversionMatrices::BT601_LIMITED;
     } else if (codecContext->colorspace == AVCOL_SPC_BT709) {
       return ColorConversionMatrices::BT709_LIMITED;
     } else if (codecContext->colorspace == AVCOL_SPC_BT2020_NCL) {
       return ColorConversionMatrices::BT2020_LIMITED;
-    } else { // AVCOL_SPC_UNSPECIFIED or unknown - default to BT.601
+    } else { // default to BT.601
       return ColorConversionMatrices::BT601_LIMITED;
     }
   } else if (codecContext->color_range == AVCOL_RANGE_JPEG) { // full range
@@ -428,7 +426,7 @@ const Npp32f (*getConversionMatrix(AVCodecContext* codecContext))[4] {
       return ColorConversionMatrices::BT709_FULL;
     } else if (codecContext->colorspace == AVCOL_SPC_BT2020_NCL) {
       return ColorConversionMatrices::BT2020_FULL;
-    } else { // AVCOL_SPC_UNSPECIFIED or unknown - default to BT.601
+    } else { // default to BT.601
       return ColorConversionMatrices::BT601_FULL;
     }
   }

src/torchcodec/_core/Encoder.cpp

@@ -787,6 +787,7 @@ void VideoEncoder::initializeEncoder(
         validatePixelFormat(*avCodec, videoStreamOptions.pixelFormat.value());
   } else {
     if (frames_.device().is_cuda()) {
+      // Default to YUV420P for CUDA encoding if unset.
       outPixelFormat_ = AV_PIX_FMT_YUV420P;
     } else {
       const AVPixelFormat* formats = getSupportedPixelFormats(*avCodec);

test/test_encoders.py

@@ -1455,17 +1455,10 @@ def test_nvenc_against_ffmpeg_cli(
             )
 
         if method == "to_file":
-            metadata_fields = ["pix_fmt", "color_range", "color_space"]
+            metadata_fields = ["color_range", "color_space"]
             ffmpeg_metadata = self._get_video_metadata(
                 ffmpeg_encoded_path, metadata_fields
             )
             encoder_metadata = self._get_video_metadata(encoder_output, metadata_fields)
             assert encoder_metadata["color_range"] == ffmpeg_metadata["color_range"]
             assert encoder_metadata["color_space"] == ffmpeg_metadata["color_space"]
-
-            encoder_fmt = encoder_metadata["pix_fmt"]
-            ffmpeg_fmt = ffmpeg_metadata["pix_fmt"]
-            # Accept yuv420p (modern) as equivalent to yuvj420p (deprecated full range)
-            assert (encoder_fmt == ffmpeg_fmt) or (
-                encoder_fmt == "yuv420p" and ffmpeg_fmt == "yuvj420p"
-            )