Split initialize_video into initialize_video_decoding and initialize_color_conversion

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -324,16 +324,21 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device)
   nvcuvid_available_ = load_nvcuvid_library();
 }
 
-void BetaCudaDeviceInterface::initialize_video(
-    const AVStream* av_stream,
-    const UniqueDecodingAVFormatContext& av_format_ctx,
+void BetaCudaDeviceInterface::initialize_color_conversion(
     const VideoStreamOptions& video_stream_options,
     const std::vector<std::unique_ptr<Transform>>& transforms,
     const std::optional<FrameDims>& resized_output_dims) {
-  // TODO_API_BREAKDOWN P0
-  if (!av_stream) {
-    return;
+  output_dtype_ = video_stream_options.output_dtype;
+  if (cpu_fallback_) {
+    cpu_fallback_->initialize_color_conversion(
+        video_stream_options, transforms, resized_output_dims);
   }
+}
+
+void BetaCudaDeviceInterface::initialize_video_decoding(
+    const AVStream* av_stream,
+    const UniqueDecodingAVFormatContext& av_format_ctx,
+    const VideoStreamOptions& video_stream_options) {
   STD_TORCH_CHECK(av_stream != nullptr, "AVStream cannot be null");
   CudaContextGuard context_guard(device_.index());
   rotation_ = rotation_from_degrees(get_rotation_from_stream(av_stream));
@@ -354,12 +359,8 @@ void BetaCudaDeviceInterface::initialize_video(
     STD_TORCH_CHECK(
         cpu_fallback_ != nullptr, "Failed to create CPU device interface");
     cpu_fallback_->initialize(codec_context_);
-    cpu_fallback_->initialize_video(
-        av_stream,
-        av_format_ctx,
-        video_stream_options,
-        transforms,
-        resized_output_dims);
+    cpu_fallback_->initialize_video_decoding(
+        av_stream, av_format_ctx, video_stream_options);
     return;
   }
 

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -45,9 +45,12 @@ class BetaCudaDeviceInterface : public DeviceInterface {
 
   void initialize(const SharedAVCodecContext& codec_context) override;
 
-  void initialize_video(
+  void initialize_video_decoding(
       const AVStream* av_stream,
       const UniqueDecodingAVFormatContext& av_format_ctx,
+      const VideoStreamOptions& video_stream_options) override;
+
+  void initialize_color_conversion(
       const VideoStreamOptions& video_stream_options,
       const std::vector<std::unique_ptr<Transform>>& transforms,
       const std::optional<FrameDims>& resized_output_dims) override;

src/torchcodec/_core/ColorConverter.cpp

@@ -28,20 +28,10 @@ ColorConverter::ColorConverter(const StableDevice& device) {
   // TODO_API_BREAKDOWN P1 It seems unnatural that the color-converter needs its
   // own device_interface_, but at the same time the color-conversion *must* be
   // third-party aware, and the only way to achieve that for now is via the
-  // interface. Should at the very least write a note about this design that now
-  // the DeviceInterface has different modes: decode only, color-convert only,
-  // and decode+color-convert (which used to be the only mode).
-
-  // TODO_API_BREAKDOWN P0: we shouldn't call initialize_video here, this is for
-  // the decoding+color-convert mode. We should do something cleaner e.g.
-  // initialize_color_convertion_only()
+  // interface.
   std::vector<std::unique_ptr<Transform>> no_transforms;
-  device_interface_->initialize_video(
-      /*av_stream=*/nullptr,
-      UniqueDecodingAVFormatContext{},
-      options,
-      no_transforms,
-      /*resized_output_dims=*/std::nullopt);
+  device_interface_->initialize_color_conversion(
+      options, no_transforms, /*resized_output_dims=*/std::nullopt);
 }
 
 torch::stable::Tensor ColorConverter::convert(const AVFrame& av_frame) {

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -80,21 +80,18 @@ void CpuDeviceInterface::initialize(const SharedAVCodecContext& codec_context) {
   codec_context_ = codec_context;
 }
 
-void CpuDeviceInterface::initialize_video(
+void CpuDeviceInterface::initialize_video_decoding(
     const AVStream* av_stream,
     [[maybe_unused]] const UniqueDecodingAVFormatContext& av_format_ctx,
+    [[maybe_unused]] const VideoStreamOptions& video_stream_options) {
+  STD_TORCH_CHECK(av_stream != nullptr, "avStream is null");
+  time_base_ = av_stream->time_base;
+}
+
+void CpuDeviceInterface::initialize_color_conversion(
     const VideoStreamOptions& video_stream_options,
     const std::vector<std::unique_ptr<Transform>>& transforms,
     const std::optional<FrameDims>& resized_output_dims) {
-  // TODO_API_BREAKDOWN P0 this used to be:
-  // STD_TORCH_CHECK(av_stream != nullptr, "avStream is null");
-  // time_base_ = av_stream->time_base;
-  // but now that avStrean can be null (to create a standalone color converter)
-  // we need this workaround. This is bad, we need to preserve the previous
-  // check somehow. See corresponding TODO in color-converter and packet decoder
-  // code.
-  time_base_ = (av_stream != nullptr) ? av_stream->time_base
-                                      : AVRational{1, AV_TIME_BASE};
   av_media_type_ = AVMEDIA_TYPE_VIDEO;
   video_stream_options_ = video_stream_options;
   resized_output_dims_ = resized_output_dims;

src/torchcodec/_core/CpuDeviceInterface.h

@@ -27,9 +27,12 @@ class CpuDeviceInterface : public DeviceInterface {
 
   virtual void initialize(const SharedAVCodecContext& codec_context) override;
 
-  virtual void initialize_video(
+  virtual void initialize_video_decoding(
       const AVStream* av_stream,
       const UniqueDecodingAVFormatContext& av_format_ctx,
+      const VideoStreamOptions& video_stream_options) override;
+
+  virtual void initialize_color_conversion(
       const VideoStreamOptions& video_stream_options,
       const std::vector<std::unique_ptr<Transform>>& transforms,
       const std::optional<FrameDims>& resized_output_dims) override;
@@ -76,7 +79,11 @@ class CpuDeviceInterface : public DeviceInterface {
       const FrameDims& output_dims) const;
 
   VideoStreamOptions video_stream_options_;
-  AVRational time_base_;
+  // Default used when color conversion runs standalone (no stream to derive it
+  // from, e.g. the ColorConverter block API). initialize_video_decoding()
+  // overrides it from the stream when there is one. Its value doesn't matter on
+  // color-conversion-only mode, but filtergraph still expects it.
+  AVRational time_base_ = {1, AV_TIME_BASE};
   AVPixelFormat output_pixel_format_;
 
   // If the resized output dimensions are present, then we always use those as

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -114,17 +114,13 @@ void CudaDeviceInterface::initialize(
   codec_context_ = codec_context;
 }
 
-void CudaDeviceInterface::initialize_video(
+void CudaDeviceInterface::initialize_video_decoding(
     const AVStream* av_stream,
     const UniqueDecodingAVFormatContext& av_format_ctx,
-    const VideoStreamOptions& video_stream_options,
-    [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>& transforms,
-    [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims) {
+    [[maybe_unused]] const VideoStreamOptions& video_stream_options) {
   STD_TORCH_CHECK(av_stream != nullptr, "avStream is null");
   time_base_ = av_stream->time_base;
-  video_stream_options_ = video_stream_options;
 
-  // TODO: Ideally, we should keep all interface implementations independent.
   cpu_interface_ = create_device_interface(kStableCPU);
   STD_TORCH_CHECK(
       cpu_interface_ != nullptr, "Failed to create CPU device interface");
@@ -137,6 +133,13 @@ void CudaDeviceInterface::initialize_video(
       /*resizedOutputDims=*/std::nullopt);
 }
 
+void CudaDeviceInterface::initialize_color_conversion(
+    const VideoStreamOptions& video_stream_options,
+    [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>& transforms,
+    [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims) {
+  video_stream_options_ = video_stream_options;
+}
+
 void CudaDeviceInterface::register_hardware_device_with_codec(
     AVCodecContext* codec_context) {
   STD_TORCH_CHECK(

src/torchcodec/_core/CudaDeviceInterface.h

@@ -28,14 +28,15 @@ class CudaDeviceInterface : public DeviceInterface {
 
   void initialize(const SharedAVCodecContext& codec_context) override;
 
-  void initialize_video(
+  void initialize_video_decoding(
       const AVStream* av_stream,
       const UniqueDecodingAVFormatContext& av_format_ctx,
+      const VideoStreamOptions& video_stream_options) override;
+
+  void initialize_color_conversion(
       const VideoStreamOptions& video_stream_options,
-      [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>&
-          transforms,
-      [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims)
-      override;
+      const std::vector<std::unique_ptr<Transform>>& transforms,
+      const std::optional<FrameDims>& resized_output_dims) override;
 
   void register_hardware_device_with_codec(
       AVCodecContext* codec_context) override;

src/torchcodec/_core/DeviceInterface.h

@@ -56,15 +56,33 @@ class DeviceInterface {
   // default sendPacket/receiveFrame/flush implementations.
   virtual void initialize(const SharedAVCodecContext& codec_context) = 0;
 
-  // Initialize the device with parameters specific to video decoding. There is
-  // a default empty implementation.
-  virtual void initialize_video(
+  // Initialize state needed to decode packets into raw AVFrames.
+  virtual void initialize_video_decoding(
       [[maybe_unused]] const AVStream* av_stream,
       [[maybe_unused]] const UniqueDecodingAVFormatContext& av_format_ctx,
+      [[maybe_unused]] const VideoStreamOptions& video_stream_options) {}
+
+  // Initialize state needed to color-convert decoded AVFrames into output
+  // tensors.
+  virtual void initialize_color_conversion(
       [[maybe_unused]] const VideoStreamOptions& video_stream_options,
       [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>&
-          transforms,
-      [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims) {}
+          transforms = {},
+      [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims =
+          std::nullopt) {}
+
+  // Convenience for the combined decode + color-convert path, kept for BC as
+  // it's used by SingleStreamDecoder and out-of-tree interfaces rely on it.
+  void initialize_video(
+      const AVStream* av_stream,
+      const UniqueDecodingAVFormatContext& av_format_ctx,
+      const VideoStreamOptions& video_stream_options,
+      const std::vector<std::unique_ptr<Transform>>& transforms,
+      const std::optional<FrameDims>& resized_output_dims) {
+    initialize_video_decoding(av_stream, av_format_ctx, video_stream_options);
+    initialize_color_conversion(
+        video_stream_options, transforms, resized_output_dims);
+  }
 
   // Initialize the device with parameters specific to audio decoding. There is
   // a default empty implementation.

src/torchcodec/_core/PacketDecoder.cpp

@@ -75,16 +75,8 @@ PacketDecoder::PacketDecoder(
   options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet
   options.device = device;
 
-  // TODO_API_BREAKDOWN P0: This isn't right, it's needed only for the NVDEC
-  // interface. This should probably be initialize_video_only - there's a
-  // sibling TODO in the ColorConverter code (about color-conversion only.)
-  std::vector<std::unique_ptr<Transform>> no_transforms;
-  device_interface_->initialize_video(
-      stream,
-      demuxer.format_context(),
-      options,
-      no_transforms,
-      /*resized_output_dims=*/std::nullopt);
+  device_interface_->initialize_video_decoding(
+      stream, demuxer.format_context(), options);
 }
 
 int PacketDecoder::send_packet(AVPacket* packet) {