Tag NVDEC surfaces with their real pixel format (#1624)

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -84,7 +84,7 @@ cudaVideoSurfaceFormat get_preferred_surface_format(OutputDtype output_dtype) {
 // when decoding happens, but this interface can be used in
 // color-conversion-only mode.
 bool is_cpu_fallback(int format) {
-  return format != AV_PIX_FMT_NV12 && format != AV_PIX_FMT_P016LE;
+  return format != AV_PIX_FMT_NV12 && !is_nvdec_16bit_surface(format);
 }
 
 static bool g_cuda_nvdec = register_device_interface(
@@ -825,9 +825,9 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame(
 
   av_frame->width = width;
   av_frame->height = height;
-  av_frame->format = (surface_format_ == cudaVideoSurfaceFormat_P016)
-      ? AV_PIX_FMT_P016LE
-      : AV_PIX_FMT_NV12;
+  av_frame->format = nvdec_pix_fmt(
+      surface_format_ == cudaVideoSurfaceFormat_P016,
+      static_cast<int>(video_format_.bit_depth_luma_minus8) + 8);
   av_frame->pts = disp_info.timestamp;
 
   // TODONVDEC P2: We compute the duration based on average frame rate info, so
@@ -1177,6 +1177,8 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
 
   UniqueAVFrame transferred_frame;
   if (cpu_fallback) {
+    // TODO: uploaded fallback frames stay tagged P016 even for 10-/12-bit
+    // sources, so they report 16 bits where an NVDEC frame reports the truth.
     AVPixelFormat target_pix_fmt = (output_dtype_ == OutputDtype::FLOAT32)
         ? AV_PIX_FMT_P016LE
         : AV_PIX_FMT_NV12;
@@ -1190,8 +1192,8 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
 
   STD_TORCH_CHECK(
       gpu_frame.format == AV_PIX_FMT_NV12 ||
-          gpu_frame.format == AV_PIX_FMT_P016LE,
-      "Expected NV12 or P016LE format frame");
+          is_nvdec_16bit_surface(gpu_frame.format),
+      "Expected NV12 or 16-bit semi-planar format frame");
 
   cudaStream_t producer_stream;
   if (mode() == Mode::ColorConverterOnly) {
@@ -1210,21 +1212,13 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
   // execrcized.
   auto convert_frame = [&](std::optional<torch::stable::Tensor> pre_alloc)
       -> torch::stable::Tensor {
-    bool is_p016 = (gpu_frame.format == AV_PIX_FMT_P016LE);
-    int bit_depth = 8;
-    if (is_p016) {
-      bit_depth = cpu_fallback
-          ? codec_context_->bits_per_raw_sample
-          : static_cast<int>(video_format_.bit_depth_luma_minus8) + 8;
-    }
     return convert_yuv_frame_to_rgb(
         gpu_frame,
         device_,
         producer_stream,
         pre_alloc,
         original_dims,
-        is_p016,
-        bit_depth,
+        static_cast<AVPixelFormat>(gpu_frame.format),
         cached_color_matrix_);
   };
 

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -345,8 +345,7 @@ void CudaDeviceInterface::convert_av_frame_to_frame_output(
       nvdec_stream,
       pre_allocated_output_tensor,
       FrameDims(av_frame.height, av_frame.width),
-      /*isP016=*/false,
-      /*bitDepth=*/8,
+      actual_format,
       cached_color_matrix_);
 }
 

src/torchcodec/_core/FFMPEGCommon.cpp

@@ -16,6 +16,45 @@ extern "C" {
 
 namespace facebook::torchcodec {
 
+#if FFMPEG_HAS_P012
+// takes is_p016_surface as input instead of the actual NVDEC surface type so we
+// don't have to include the NVDEC headers here
+AVPixelFormat nvdec_pix_fmt(bool is_p016_surface, int bit_depth) {
+  if (!is_p016_surface) {
+    return AV_PIX_FMT_NV12;
+  }
+  switch (bit_depth) {
+    case 10:
+      return AV_PIX_FMT_P010LE;
+    case 12:
+      return AV_PIX_FMT_P012LE;
+    default:
+      return AV_PIX_FMT_P016LE;
+  }
+}
+#else
+AVPixelFormat nvdec_pix_fmt(bool is_p016_surface, int bit_depth) {
+  // TODO_API_BREAKDOWN P2: needs a comment about P012 missing and why it's
+  // still OK to return P016LE.
+  if (!is_p016_surface) {
+    return AV_PIX_FMT_NV12;
+  }
+  return bit_depth == 10 ? AV_PIX_FMT_P010LE : AV_PIX_FMT_P016LE;
+}
+#endif // FFMPEG_HAS_P012
+
+#if FFMPEG_HAS_P012
+bool is_nvdec_16bit_surface(int format) {
+  return format == AV_PIX_FMT_P010LE || format == AV_PIX_FMT_P012LE ||
+      format == AV_PIX_FMT_P016LE;
+}
+#else
+
+bool is_nvdec_16bit_surface(int format) {
+  return format == AV_PIX_FMT_P010LE || format == AV_PIX_FMT_P016LE;
+}
+#endif
+
 AutoAVPacket::AutoAVPacket() : av_packet_(av_packet_alloc()) {
   STD_TORCH_CHECK(av_packet_ != nullptr, "Couldn't allocate avPacket.");
 }

src/torchcodec/_core/FFMPEGCommon.h

@@ -44,6 +44,13 @@ extern "C" {
 #define FFMPEG_HAS_CH_LAYOUT 0
 #endif
 
+// FFmpeg 6 added AV_PIX_FMT_P012, the 12-bit sibling of P010.
+#if LIBAVUTIL_VERSION_MAJOR >= 58
+#define FFMPEG_HAS_P012 1
+#else
+#define FFMPEG_HAS_P012 0
+#endif
+
 // FFmpeg 7.1 added avcodec_get_supported_config(), replacing the codec's
 // pix_fmts / sample_fmts / supported_samplerates / ch_layouts arrays.
 #if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(61, 13, 100)
@@ -61,6 +68,9 @@ extern "C" {
 
 namespace facebook::torchcodec {
 
+AVPixelFormat nvdec_pix_fmt(bool is_p016_surface, int bit_depth);
+bool is_nvdec_16bit_surface(int format);
+
 // FFMPEG uses special delete functions for some structures. These template
 // functions are used to pass into unique_ptr as custom deleters so we can
 // wrap FFMPEG structs with unique_ptrs for ease of use.

src/torchcodec/_core/PacketDecoder.cpp

@@ -140,8 +140,7 @@ FramePlanes frame_to_planes(
     int64_t bytes_per_sample = (comp.depth > 8) ? 2 : 1;
     int64_t linesize = av_frame.linesize[comp.plane];
     STD_TORCH_CHECK(
-        comp.shift == 0 && comp.depth <= 16 && linesize > 0 &&
-            comp.step % bytes_per_sample == 0 &&
+        comp.depth <= 16 && linesize > 0 && comp.step % bytes_per_sample == 0 &&
             linesize % bytes_per_sample == 0,
         "Cannot expose component ",
         c,

src/torchcodec/_core/color_conversion.cpp

@@ -194,9 +194,13 @@ torch::stable::Tensor convert_yuv_frame_to_rgb(
     cudaStream_t producer_stream,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor,
     const FrameDims& output_dims,
-    bool is_p016,
-    int bit_depth,
+    AVPixelFormat pix_fmt,
     CachedColorMatrix& cached_color_matrix) {
+  bool is_p016 = is_nvdec_16bit_surface(pix_fmt);
+  const AVPixFmtDescriptor* desc = av_pix_fmt_desc_get(pix_fmt);
+  STD_TORCH_CHECK(desc != nullptr, "Unknown pixel format on decoded frame");
+  int bit_depth = desc->comp[0].depth;
+
   float out_scale = is_p016 ? 65535.0f : 255.0f;
   OutputDtype out_dtype = is_p016 ? OutputDtype::FLOAT32 : OutputDtype::UINT8;
 

src/torchcodec/_core/color_conversion.h

@@ -79,18 +79,16 @@ void launch_p016_to_rgb16_kernel(
 
 // Convert a YUV frame (NV12 or P016) on GPU to an interleaved RGB tensor.
 //
-// isP016: true for P016 (uint16 I/O), false for NV12 (uint8 I/O).
-// bitDepth: 8 for NV12, actual bit depth (10 or 12) for P016.
 // outputDims: desired output size; if the frame was rounded up to even
 //   dimensions, the result is cropped back to outputDims.
+// pixFmt: the format the samples are actually in (NV12, P010, P012, P016).
 torch::stable::Tensor convert_yuv_frame_to_rgb(
     const AVFrame& av_frame,
     const StableDevice& device,
     cudaStream_t producer_stream,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor,
     const FrameDims& output_dims,
-    bool is_p016,
-    int bit_depth,
+    AVPixelFormat pix_fmt,
     CachedColorMatrix& cached_color_matrix);
 
 // Compute the RGB -> YUV color conversion matrix (for encoding).