Cover the frames that aren't three YUV planes

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -4,6 +4,7 @@
 // This source code is licensed under the BSD-style license found in the
 // LICENSE file in the root directory of this source tree.
 
+#include <limits>
 #include <map>
 #include <mutex>
 #include <vector>
@@ -343,7 +344,7 @@ std::optional<cudaVideoSurfaceFormat> get_nvdec_surface_format(
 void standalone_frame_free_callback(
     [[maybe_unused]] void* opaque,
     uint8_t* data) {
-  delete reinterpret_cast<StandAloneFrameAttachedData*>(data);
+  delete reinterpret_cast<OwnedFrameStorage*>(data);
 }
 
 class CudaContextGuard {
@@ -964,7 +965,12 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame(
   av_frame->data[1] = plane(1);
   av_frame->data[2] = is_444 ? plane(2) : nullptr;
   av_frame->data[3] = nullptr;
-  // TODO_API_BREAKDOWN CC P2: Check range before cast?
+  STD_TORCH_CHECK(
+      pitch <= static_cast<unsigned int>(std::numeric_limits<int>::max()),
+      "NVDEC returned a pitch of ",
+      pitch,
+      " bytes, which doesn't fit in an AVFrame line size. This should never "
+      "happen, please report.");
   av_frame->linesize[0] = static_cast<int>(pitch);
   av_frame->linesize[1] = static_cast<int>(pitch);
   av_frame->linesize[2] = is_444 ? static_cast<int>(pitch) : 0;
@@ -1000,12 +1006,12 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) {
     storage = copy_nvdec_surface(av_frame, current_stream);
   }
 
-  auto attached_data = new StandAloneFrameAttachedData();
+  auto attached_data = new OwnedFrameStorage();
   attached_data->frame_ready.record(current_stream);
   attached_data->storage = std::move(storage);
   av_frame->opaque_ref = av_buffer_create(
       reinterpret_cast<uint8_t*>(attached_data),
-      sizeof(StandAloneFrameAttachedData),
+      sizeof(OwnedFrameStorage),
       standalone_frame_free_callback,
       nullptr,
       0);
@@ -1065,8 +1071,7 @@ std::optional<torch::stable::Tensor> BetaCudaDeviceInterface::get_frame_storage(
   // for those users who would like to consume the frame with their own
   // consumer, i.e. not using the ColorConverter: they need to call
   // frame.storage.record_stream(color_conversion_stream) themselves.
-  return reinterpret_cast<StandAloneFrameAttachedData*>(
-             av_frame.opaque_ref->data)
+  return reinterpret_cast<OwnedFrameStorage*>(av_frame.opaque_ref->data)
       ->storage;
 }
 
@@ -1233,10 +1238,21 @@ GpuFrameAndStorage BetaCudaDeviceInterface::upload_cpu_frame_to_gpu(
       cpu_frame.colorspace,
       width,
       height,
-      target_pix_fmt);
+      target_pix_fmt,
+      // The frame keeps its range tag through the upload, so the samples must
+      // keep the range that tag names. Left to itself, swscale writes limited
+      // range into a YUV target whatever it read, and the color conversion
+      // would then expand a full-range source a second time.
+      cpu_frame.color_range);
 
   if (!sws_context_ || prev_sws_config_ != sws_config) {
-    sws_context_ = create_sws_context(sws_config, SWS_BILINEAR);
+    // Nothing is rescaled here, so the flags only pick how chroma is
+    // resampled, which happens when the source is subsampled more finely than
+    // the target surface (4:2:2 into 4:4:4, say). SWS_POINT replicates it,
+    // which is what the CPU converter does on its way to RGB - interpolating
+    // instead would invent chroma the CPU never sees, and show up as colored
+    // fringes along sharp edges.
+    sws_context_ = create_sws_context(sws_config, SWS_POINT);
     prev_sws_config_ = sws_config;
   }
 
@@ -1323,6 +1339,13 @@ GpuFrameAndStorage BetaCudaDeviceInterface::upload_cpu_frame_to_gpu(
       "Failed to copy frame properties: ",
       get_ffmpeg_error_string_from_error_code(ret));
 
+  // AVCOL_SPC_RGB says "these planes are RGB", which the planes we just wrote
+  // aren't. Name the matrix swscale encoded them with instead: it maps
+  // AVCOL_SPC_RGB, like any colorspace it doesn't know, to its BT.601 default.
+  if (cpu_frame.colorspace == AVCOL_SPC_RGB) {
+    gpu_frame->colorspace = AVCOL_SPC_SMPTE170M;
+  }
+
   return {std::move(gpu_frame), std::move(storage)};
 }
 
@@ -1369,8 +1392,8 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
         gpu_frame.opaque_ref != nullptr,
         "ColorConverter received a non-standalone frame; frames fed to a "
         "standalone ColorConverter must come from a PacketDecoder.");
-    auto attached_data = reinterpret_cast<StandAloneFrameAttachedData*>(
-        gpu_frame.opaque_ref->data);
+    auto attached_data =
+        reinterpret_cast<OwnedFrameStorage*>(gpu_frame.opaque_ref->data);
     attached_data->frame_ready.make_stream_wait(current_stream);
   } else {
     STD_TORCH_CHECK(

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -34,9 +34,8 @@
 #include "nvcuvid_include/nvcuvid.h"
 
 namespace facebook::torchcodec {
-// TODO_API_BREAKDOWN P2: the name says "standalone", but this is really about
-// owning a GPU buffer. Find one that covers both.
-struct StandAloneFrameAttachedData {
+// The buffer a frame owns its samples in, hung off the AVFrame as opaque data.
+struct OwnedFrameStorage {
   // Marks the point where the copy (or upload) that filled `storage` was
   // enqueued. A consumer on another stream must wait on it.
   CudaEvent frame_ready;

src/torchcodec/_core/FFMPEGCommon.cpp

@@ -904,14 +904,16 @@ SwsConfig::SwsConfig(
     AVColorSpace input_colorspace,
     int output_width,
     int output_height,
-    AVPixelFormat output_format)
+    AVPixelFormat output_format,
+    AVColorRange output_color_range)
     : input_width(input_width),
       input_height(input_height),
       input_format(input_format),
       input_colorspace(input_colorspace),
       output_width(output_width),
       output_height(output_height),
-      output_format(output_format) {}
+      output_format(output_format),
+      output_color_range(output_color_range) {}
 
 bool SwsConfig::operator==(const SwsConfig& other) const {
   return input_width == other.input_width &&
@@ -920,7 +922,8 @@ bool SwsConfig::operator==(const SwsConfig& other) const {
       input_colorspace == other.input_colorspace &&
       output_width == other.output_width &&
       output_height == other.output_height &&
-      output_format == other.output_format;
+      output_format == other.output_format &&
+      output_color_range == other.output_color_range;
 }
 
 bool SwsConfig::operator!=(const SwsConfig& other) const {
@@ -957,6 +960,13 @@ UniqueSwsContext create_sws_context(
       &saturation);
   STD_TORCH_CHECK(ret != -1, "sws_getColorspaceDetails returned -1");
 
+  // swscale spells a range as an int: 1 is full, 0 is limited. FFmpeg names
+  // those AVCOL_RANGE_JPEG and AVCOL_RANGE_MPEG, after the two worlds they come
+  // from - JPEG is the full one.
+  if (sws_config.output_color_range != AVCOL_RANGE_UNSPECIFIED) {
+    dst_range = sws_config.output_color_range == AVCOL_RANGE_JPEG;
+  }
+
   const int* colorspace_table =
       sws_getCoefficients(sws_config.input_colorspace);
   ret = sws_setColorspaceDetails(

src/torchcodec/_core/FFMPEGCommon.h

@@ -406,6 +406,10 @@ struct SwsConfig {
   int output_width = 0;
   int output_height = 0;
   AVPixelFormat output_format = AV_PIX_FMT_NONE;
+  // swscale derives the output range from the output pixel format alone: YUV
+  // gets limited range. AVCOL_RANGE_UNSPECIFIED keeps that; anything else
+  // overrides it.
+  AVColorRange output_color_range = AVCOL_RANGE_UNSPECIFIED;
 
   SwsConfig() = default;
   SwsConfig(
@@ -415,7 +419,8 @@ struct SwsConfig {
       AVColorSpace input_colorspace,
       int output_width,
       int output_height,
-      AVPixelFormat output_format);
+      AVPixelFormat output_format,
+      AVColorRange output_color_range = AVCOL_RANGE_UNSPECIFIED);
 
   bool operator==(const SwsConfig& other) const;
   bool operator!=(const SwsConfig& other) const;