HDR-correct dtype handling for the Blocks API (#1625)

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -1208,8 +1208,6 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output(
     producer_stream = get_current_cuda_stream(device_.index());
   }
 
-  // TODO_API_BREAKDOWN P1: we don't suppor output_dtype so some of that is not
-  // execrcized.
   auto convert_frame = [&](std::optional<torch::stable::Tensor> pre_alloc)
       -> torch::stable::Tensor {
     return convert_yuv_frame_to_rgb(

src/torchcodec/_core/ColorConverter.cpp

@@ -15,15 +15,28 @@
 
 namespace facebook::torchcodec {
 
-ColorConverter::ColorConverter(const StableDevice& device) {
+ColorConverter::ColorConverter(
+    const StableDevice& device,
+    OutputDtypeConfig output_dtype_config)
+    : device_(device), output_dtype_config_(output_dtype_config) {
   device_interface_ = create_device_interface(device);
   STD_TORCH_CHECK(
       device_interface_ != nullptr,
       "Failed to create device interface. This should never happen, please report.");
+}
+
+void ColorConverter::maybe_initialize_interface(OutputDtype output_dtype) {
+  // Interface initialization is done per-frame, not in the constructor: with
+  // AUTO, the desired output dtype is only known once we see a frame, and it
+  // can differ from one frame to the next.
+  if (initialized_output_dtype_.has_value() &&
+      *initialized_output_dtype_ == output_dtype) {
+    return;
+  }
 
   VideoStreamOptions options;
-  options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet
-  options.device = device;
+  options.output_dtype = output_dtype;
+  options.device = 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
@@ -32,13 +45,18 @@ ColorConverter::ColorConverter(const StableDevice& device) {
   std::vector<std::unique_ptr<Transform>> no_transforms;
   device_interface_->initialize_color_conversion(
       options, no_transforms, /*resized_output_dims=*/std::nullopt);
+  initialized_output_dtype_ = output_dtype;
 }
 
 torch::stable::Tensor ColorConverter::convert(const AVFrame& av_frame) {
+  OutputDtype output_dtype = resolve_output_dtype(
+      output_dtype_config_, static_cast<AVPixelFormat>(av_frame.format));
+  maybe_initialize_interface(output_dtype);
+
   FrameOutput frame_output;
   device_interface_->convert_av_frame_to_frame_output(
       av_frame, frame_output, std::nullopt);
-  return frame_output.data;
+  return convert_to_output_dtype(frame_output.data, output_dtype);
 }
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/ColorConverter.h

@@ -7,23 +7,31 @@
 #pragma once
 
 #include <memory>
+#include <optional>
 #include <string_view>
 
 #include "DeviceInterface.h"
 #include "FFMPEGCommon.h"
 #include "StableABICompat.h"
+#include "StreamOptions.h"
 
 namespace facebook::torchcodec {
 
 class FORCE_PUBLIC_VISIBILITY ColorConverter {
  public:
   explicit ColorConverter(
-      const StableDevice& device = StableDevice(kStableCPU));
+      const StableDevice& device = StableDevice(kStableCPU),
+      OutputDtypeConfig output_dtype_config = OutputDtypeConfig::UINT8);
 
   torch::stable::Tensor convert(const AVFrame& av_frame);
 
  private:
+  void maybe_initialize_interface(OutputDtype output_dtype);
+
   std::unique_ptr<DeviceInterface> device_interface_;
+  StableDevice device_;
+  OutputDtypeConfig output_dtype_config_;
+  std::optional<OutputDtype> initialized_output_dtype_;
 };
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/FFMPEGCommon.cpp

@@ -55,6 +55,25 @@ bool is_nvdec_16bit_surface(int format) {
 }
 #endif
 
+OutputDtype resolve_output_dtype(
+    OutputDtypeConfig output_dtype_config,
+    AVPixelFormat pix_fmt) {
+  switch (output_dtype_config) {
+    case OutputDtypeConfig::UINT8:
+      return OutputDtype::UINT8;
+    case OutputDtypeConfig::FLOAT32:
+      return OutputDtype::FLOAT32;
+    case OutputDtypeConfig::AUTO: {
+      // TODO_HDR: This is basically our heuristic that defines how we identify
+      // HDR videos, we might want to refine it.
+      const AVPixFmtDescriptor* desc = av_pix_fmt_desc_get(pix_fmt);
+      return (desc != nullptr && desc->comp[0].depth > 8) ? OutputDtype::FLOAT32
+                                                          : OutputDtype::UINT8;
+    }
+  }
+  return OutputDtype::UINT8;
+}
+
 AutoAVPacket::AutoAVPacket() : av_packet_(av_packet_alloc()) {
   STD_TORCH_CHECK(av_packet_ != nullptr, "Couldn't allocate avPacket.");
 }

src/torchcodec/_core/FFMPEGCommon.h

@@ -12,6 +12,8 @@
 #include <string>
 #include <vector>
 
+#include "StreamOptions.h"
+
 extern "C" {
 #include <libavcodec/avcodec.h>
 #include <libavcodec/bsf.h>
@@ -71,6 +73,10 @@ namespace facebook::torchcodec {
 AVPixelFormat nvdec_pix_fmt(bool is_p016_surface, int bit_depth);
 bool is_nvdec_16bit_surface(int format);
 
+OutputDtype resolve_output_dtype(
+    OutputDtypeConfig output_dtype_config,
+    AVPixelFormat pix_fmt);
+
 // 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/Frame.cpp

@@ -53,4 +53,23 @@ torch::stable::Tensor allocate_empty_hwc_tensor(
   }
 }
 
+torch::stable::Tensor convert_to_output_dtype(
+    const torch::stable::Tensor& tensor,
+    OutputDtype output_dtype) {
+  bool is_uint16 = tensor.scalar_type() == kStableUInt16;
+
+  if (output_dtype == OutputDtype::FLOAT32) {
+    double max_val = is_uint16 ? 65535.0 : 255.0;
+    return stable_div(torch::stable::to(tensor, kStableFloat32), max_val);
+  }
+
+  if (!is_uint16) {
+    return tensor;
+  }
+  // uint16 -> uint8. 257 is 65535 / 255.
+  return torch::stable::to(
+      stable_div(torch::stable::to(tensor, kStableFloat32), 257.0),
+      kStableUInt8);
+}
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/Frame.h

@@ -69,4 +69,8 @@ torch::stable::Tensor allocate_empty_hwc_tensor(
     OutputDtype output_dtype,
     std::optional<int> num_frames = std::nullopt);
 
+torch::stable::Tensor convert_to_output_dtype(
+    const torch::stable::Tensor& tensor,
+    OutputDtype output_dtype);
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/PacketDecoder.cpp

@@ -71,11 +71,17 @@ PacketDecoder::PacketDecoder(
       stream, av_codec, device_interface_.get(), ffmpeg_thread_count);
   device_interface_->initialize(codec_context_);
 
+  const AVPixFmtDescriptor* stream_desc =
+      av_pix_fmt_desc_get(codec_context_->pix_fmt);
+  int stream_bit_depth = stream_desc ? stream_desc->comp[0].depth : 8;
+
   VideoStreamOptions options;
-  // TODO_API_BREAKDOWN P1: Need to design and figure out behavior of Blocks
-  // with HDR data.
-  options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet
   options.device = device;
+  // This is ugly: what we actually mean is "let the device interface decode
+  // into the native surface", which matters for NVDEC.
+  // TODO_API_BREAKDOWN P2: Find a cleaner way to express this?
+  options.output_dtype =
+      stream_bit_depth > 8 ? OutputDtype::FLOAT32 : OutputDtype::UINT8;
 
   device_interface_->initialize_video_decoding(
       stream, demuxer.format_context(), options);
@@ -135,6 +141,8 @@ FramePlanes frame_to_planes(
   result.colorspace = colorspace_name ? colorspace_name : "unknown";
   result.color_range = color_range_name ? color_range_name : "unknown";
 
+  result.bit_depth = desc->comp[0].depth;
+
   for (int c = 0; c < desc->nb_components; ++c) {
     const AVComponentDescriptor& comp = desc->comp[c];
     int64_t bytes_per_sample = (comp.depth > 8) ? 2 : 1;

src/torchcodec/_core/PacketDecoder.h

@@ -73,6 +73,7 @@ struct FramePlanes {
   std::string pix_fmt;
   std::string colorspace;
   std::string color_range;
+  int64_t bit_depth = 8;
 };
 
 FORCE_PUBLIC_VISIBILITY FramePlanes frame_to_planes(

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -559,26 +559,9 @@ void SingleStreamDecoder::add_video_stream(
         active_stream_index_, custom_frame_mappings.value());
   }
 
-  // Resolve the user-facing OutputDtypeConfig (which may be AUTO) into an
-  // OutputDtype that downstream code can use directly.
-  // TODO_HDR: This is basically our heuristic that defines how we identify HDR
-  // videos, we might want to refine it.
-  switch (stream_info.video_stream_options.output_dtype_config) {
-    case OutputDtypeConfig::UINT8:
-      stream_info.video_stream_options.output_dtype = OutputDtype::UINT8;
-      break;
-    case OutputDtypeConfig::FLOAT32:
-      stream_info.video_stream_options.output_dtype = OutputDtype::FLOAT32;
-      break;
-    case OutputDtypeConfig::AUTO: {
-      const AVPixFmtDescriptor* desc = av_pix_fmt_desc_get(
-          static_cast<AVPixelFormat>(stream_info.stream->codecpar->format));
-      stream_info.video_stream_options.output_dtype =
-          (desc != nullptr && desc->comp[0].depth > 8) ? OutputDtype::FLOAT32
-                                                       : OutputDtype::UINT8;
-      break;
-    }
-  }
+  stream_info.video_stream_options.output_dtype = resolve_output_dtype(
+      stream_info.video_stream_options.output_dtype_config,
+      static_cast<AVPixelFormat>(stream_info.stream->codecpar->format));
 
   // Set preRotationDims_ for the active stream. These are the raw encoded
   // dimensions from FFmpeg, used as a fallback for tensor pre-allocation when
@@ -678,7 +661,7 @@ void SingleStreamDecoder::add_audio_stream(
 FrameOutput SingleStreamDecoder::get_next_frame() {
   auto output = get_next_frame_internal();
   if (stream_infos_[active_stream_index_].av_media_type == AVMEDIA_TYPE_VIDEO) {
-    output.data = maybe_permute_and_convert_to_float32(output.data);
+    output.data = maybe_permute_and_convert_dtype(output.data);
   }
   return output;
 }
@@ -695,7 +678,7 @@ FrameOutput SingleStreamDecoder::get_next_frame_internal(
 
 FrameOutput SingleStreamDecoder::get_frame_at_index(int64_t frame_index) {
   auto frame_output = get_frame_at_index_internal(frame_index);
-  frame_output.data = maybe_permute_and_convert_to_float32(frame_output.data);
+  frame_output.data = maybe_permute_and_convert_dtype(frame_output.data);
   return frame_output;
 }
 
@@ -805,7 +788,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_at_indices(
     previous_index_in_video = index_in_video;
   }
   frame_batch_output.data =
-      maybe_permute_and_convert_to_float32(frame_batch_output.data);
+      maybe_permute_and_convert_dtype(frame_batch_output.data);
   return frame_batch_output;
 }
 
@@ -855,7 +838,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_in_range(
     frame_batch_output_duration_seconds[f] = frame_output.duration_seconds;
   }
   frame_batch_output.data =
-      maybe_permute_and_convert_to_float32(frame_batch_output.data);
+      maybe_permute_and_convert_dtype(frame_batch_output.data);
   return frame_batch_output;
 }
 
@@ -897,7 +880,7 @@ FrameOutput SingleStreamDecoder::get_frame_played_at(double seconds) {
 
   // Convert the frame to tensor.
   FrameOutput frame_output = convert_av_frame_to_frame_output(*av_frame);
-  frame_output.data = maybe_permute_and_convert_to_float32(frame_output.data);
+  frame_output.data = maybe_permute_and_convert_dtype(frame_output.data);
   return frame_output;
 }
 
@@ -987,7 +970,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_played_in_range(
         device_interface_->get_pre_allocation_dtype(
             video_stream_options.output_dtype));
     frame_batch_output.data =
-        maybe_permute_and_convert_to_float32(frame_batch_output.data);
+        maybe_permute_and_convert_dtype(frame_batch_output.data);
     return frame_batch_output;
   }
 
@@ -1061,7 +1044,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_played_in_range(
     }
 
     frame_batch_output.data =
-        maybe_permute_and_convert_to_float32(frame_batch_output.data);
+        maybe_permute_and_convert_dtype(frame_batch_output.data);
     return frame_batch_output;
   } else {
     // Note that we look at nextPts for a frame, and not its pts or duration.
@@ -1098,7 +1081,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_played_in_range(
       frame_batch_output_duration_seconds[f] = frame_output.duration_seconds;
     }
     frame_batch_output.data =
-        maybe_permute_and_convert_to_float32(frame_batch_output.data);
+        maybe_permute_and_convert_dtype(frame_batch_output.data);
 
     return frame_batch_output;
   }
@@ -1692,7 +1675,7 @@ FrameOutput SingleStreamDecoder::convert_av_frame_to_frame_output(
 // OUTPUT ALLOCATION AND SHAPE CONVERSION
 // --------------------------------------------------------------------------
 
-torch::stable::Tensor SingleStreamDecoder::maybe_permute_and_convert_to_float32(
+torch::stable::Tensor SingleStreamDecoder::maybe_permute_and_convert_dtype(
     torch::stable::Tensor& hwc_tensor) {
   // Permute HWC to CHW if needed. Returns a view of the input tensor, the
   // leading batch-dimension [N] is optional i.e. the input tensor can be 3D or
@@ -1718,19 +1701,9 @@ torch::stable::Tensor SingleStreamDecoder::maybe_permute_and_convert_to_float32(
     }
   }
 
-  // Convert to float32 and normalize to [0, 1] if needed.
-  OutputDtype output_dtype =
-      stream_infos_[active_stream_index_].video_stream_options.output_dtype;
-  if (output_dtype != OutputDtype::FLOAT32) {
-    return tensor;
-  }
-  bool is_uint16 =
-      tensor.scalar_type() == torch::headeronly::ScalarType::UInt16;
-  double max_val = static_cast<double>(
-      is_uint16 ? std::numeric_limits<uint16_t>::max()
-                : std::numeric_limits<uint8_t>::max());
-  auto as_float = torch::stable::to(tensor, kStableFloat32);
-  return stable_div(as_float, max_val);
+  return convert_to_output_dtype(
+      tensor,
+      stream_infos_[active_stream_index_].video_stream_options.output_dtype);
 }
 
 // --------------------------------------------------------------------------