Programmatically enforce resolution of OuputDtype (#1444)

Author: NicolasHug

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -578,17 +578,25 @@ void SingleStreamDecoder::addVideoStream(
         activeStreamIndex_, customFrameMappings.value());
   }
 
-  // Resolve AUTO once based on the source bit depth, so downstream code
-  // only has to handle UINT8 / FLOAT32.
-  //
-  // TODO: programmatically enforce that OutputDtype is "resolved" (i.e. not
-  // AUTO) past this point.
-  if (streamInfo.videoStreamOptions.outputDtype == OutputDtype::AUTO) {
-    const AVPixFmtDescriptor* desc = av_pix_fmt_desc_get(
-        static_cast<AVPixelFormat>(streamInfo.stream->codecpar->format));
-    streamInfo.videoStreamOptions.outputDtype =
-        (desc != nullptr && desc->comp[0].depth > 8) ? OutputDtype::FLOAT32
-                                                     : OutputDtype::UINT8;
+  // 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 (streamInfo.videoStreamOptions.outputDtypeConfig) {
+    case OutputDtypeConfig::UINT8:
+      streamInfo.videoStreamOptions.outputDtype = OutputDtype::UINT8;
+      break;
+    case OutputDtypeConfig::FLOAT32:
+      streamInfo.videoStreamOptions.outputDtype = OutputDtype::FLOAT32;
+      break;
+    case OutputDtypeConfig::AUTO: {
+      const AVPixFmtDescriptor* desc = av_pix_fmt_desc_get(
+          static_cast<AVPixelFormat>(streamInfo.stream->codecpar->format));
+      streamInfo.videoStreamOptions.outputDtype =
+          (desc != nullptr && desc->comp[0].depth > 8) ? OutputDtype::FLOAT32
+                                                       : OutputDtype::UINT8;
+      break;
+    }
   }
 
   // Set preRotationDims_ for the active stream. These are the raw encoded
@@ -1515,10 +1523,9 @@ torch::stable::Tensor SingleStreamDecoder::maybePermuteHWC2CHW(
   }
 }
 
+// TODO_HDR: should this be a single call along with maybePermuteHWC2CHW?
 torch::stable::Tensor SingleStreamDecoder::maybeConvertToFloat32(
     torch::stable::Tensor& tensor) {
-  // AUTO has been resolved to UINT8 or FLOAT32 in addVideoStream, so we only
-  // need to handle FLOAT32 here.
   OutputDtype outputDtype =
       streamInfos_[activeStreamIndex_].videoStreamOptions.outputDtype;
   if (outputDtype != OutputDtype::FLOAT32) {

src/torchcodec/_core/SingleStreamDecoder.h

@@ -277,8 +277,7 @@ class FORCE_PUBLIC_VISIBILITY SingleStreamDecoder {
   torch::stable::Tensor maybePermuteHWC2CHW(torch::stable::Tensor& hwcTensor);
 
   // Converts the tensor to float32 and normalizes to [0, 1] when the active
-  // stream's outputDtype calls for it (FLOAT32 always, or AUTO when the tensor
-  // is uint16). Otherwise returns the input unchanged.
+  // stream's outputDtype is FLOAT32. Otherwise returns the input unchanged.
   torch::stable::Tensor maybeConvertToFloat32(torch::stable::Tensor& tensor);
 
   FrameOutput convertAVFrameToFrameOutput(

src/torchcodec/_core/StreamOptions.h

@@ -21,17 +21,20 @@ enum ColorConversionLibrary {
   SWSCALE
 };
 
-// Controls the dtype of decoded frame tensors.
+// The resolved output dtype. Only UINT8 or FLOAT32 — no AUTO.
+// All code downstream of addVideoStream() should use this.
+enum class OutputDtype { UINT8, FLOAT32 };
+
+// The user-facing output dtype config, which may include AUTO.
+// AUTO is resolved in addVideoStream() into an OutputDtype.
 // UINT8: Always output uint8 tensors (default, backward compatible). Uses an
 //        8-bit / RGB24 intermediate.
 // FLOAT32: Always output float32 tensors normalized to [0, 1]. Uses a 16-bit /
 //          RGB48 intermediate so the YUV->RGB matrix output is preserved at
 //          full precision through the float cast, regardless of source bit
 //          depth.
 // AUTO: Output uint8 for SDR (<=8-bit) sources, float32 for HDR (>8-bit).
-//       Resolved upstream in addVideoStream so downstream code only ever sees
-//       UINT8 / FLOAT32.
-enum class OutputDtype { UINT8, FLOAT32, AUTO };
+enum class OutputDtypeConfig { UINT8, FLOAT32, AUTO };
 
 struct VideoStreamOptions {
   VideoStreamOptions() {}
@@ -59,9 +62,12 @@ struct VideoStreamOptions {
   // Device variant (e.g., "nvdec", "ffmpeg")
   std::string_view deviceVariant = "default";
 
-  // Controls the dtype of decoded frame tensors. Default UINT8 preserves
-  // existing behavior; FLOAT32 and AUTO are used for high-bit-depth output
-  // (e.g. HDR).
+  // The user-specified output dtype config. May be AUTO, which gets resolved
+  // in addVideoStream() into outputDtype below.
+  OutputDtypeConfig outputDtypeConfig = OutputDtypeConfig::UINT8;
+
+  // Set by addVideoStream() after resolving AUTO. All downstream code should
+  // read this field.
   OutputDtype outputDtype = OutputDtype::UINT8;
 
   // Encoding options

src/torchcodec/_core/custom_ops.cpp

@@ -516,18 +516,27 @@ void _add_video_stream(
     std::optional<torch::stable::Tensor>
         custom_frame_mappings_keyframe_indices = std::nullopt,
     std::optional<std::string> color_conversion_library = std::nullopt,
+    // TODO_HDR: see other TODO, should default be uint8 instead? Maybe it
+    // shoudn't even be optional?
+    // videoStreamOptions.outputDtype defaults to UINT8 so if
+    // output_dtype is nullopt, the videoStreamOptions will default to UINT8.
+    // But that's pretty implicit and suggests maybe this shouldn't be optional
+    // at all.
+    // TODO_UINT8 Also this currently takes strings but surely the public API
+    // will want to support torch.dtype object, we should figure out when to do
+    // the conversion.
     std::optional<std::string> output_dtype = std::nullopt) {
   VideoStreamOptions videoStreamOptions;
   videoStreamOptions.ffmpegThreadCount = num_threads;
 
   if (output_dtype.has_value()) {
     const std::string& val = *output_dtype;
     if (val == "uint8") {
-      videoStreamOptions.outputDtype = OutputDtype::UINT8;
+      videoStreamOptions.outputDtypeConfig = OutputDtypeConfig::UINT8;
     } else if (val == "float32") {
-      videoStreamOptions.outputDtype = OutputDtype::FLOAT32;
+      videoStreamOptions.outputDtypeConfig = OutputDtypeConfig::FLOAT32;
     } else if (val == "auto") {
-      videoStreamOptions.outputDtype = OutputDtype::AUTO;
+      videoStreamOptions.outputDtypeConfig = OutputDtypeConfig::AUTO;
     } else {
       STD_TORCH_CHECK(
           false,

src/torchcodec/_core/ops.py

@@ -84,6 +84,7 @@ def add_video_stream(
     custom_frame_mappings: (
         tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None
     ) = None,
+    # TODO_HDR: should default be None or uint8??
     output_dtype: str | None = None,
 ) -> None:
     custom_frame_mappings_pts: torch.Tensor | None = None