move HDR tests to VideoDecoder (#1450)

Author: NicolasHug

src/torchcodec/_core/_decoder_utils.py

@@ -166,6 +166,7 @@ def create_video_decoder(
     device_variant: str = "default",
     transforms: Sequence[DecoderTransform | nn.Module] | None = None,
     custom_frame_mappings: tuple[Tensor, Tensor, Tensor] | None = None,
+    output_dtype: str = "uint8",
 ) -> tuple[Tensor, int, VideoStreamMetadata]:
 
     decoder = create_decoder(source=source, seek_mode=seek_mode)
@@ -191,6 +192,7 @@ def create_video_decoder(
         device_variant=device_variant,
         transform_specs=transform_specs,
         custom_frame_mappings=custom_frame_mappings,
+        output_dtype=output_dtype,
     )
 
     return (decoder, stream_index, metadata)

src/torchcodec/_core/ops.py

@@ -84,8 +84,8 @@ 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,
+    # TODO_HDR
+    output_dtype: str = "uint8",
 ) -> None:
     custom_frame_mappings_pts: torch.Tensor | None = None
     custom_frame_mappings_keyframe_indices: torch.Tensor | None = None
@@ -430,7 +430,7 @@ def _add_video_stream_abstract(
     custom_frame_mappings_duration: torch.Tensor | None = None,
     custom_frame_mappings_keyframe_indices: torch.Tensor | None = None,
     color_conversion_library: str | None = None,
-    output_dtype: str | None = None,
+    output_dtype: str = "uint8",
 ) -> None:
     return
 
@@ -448,7 +448,7 @@ def add_video_stream_abstract(
     custom_frame_mappings_pts: torch.Tensor | None = None,
     custom_frame_mappings_duration: torch.Tensor | None = None,
     custom_frame_mappings_keyframe_indices: torch.Tensor | None = None,
-    output_dtype: str | None = None,
+    output_dtype: str = "uint8",
 ) -> None:
     return
 

src/torchcodec/decoders/_video_decoder.py

@@ -125,6 +125,7 @@ class VideoDecoder:
             :class:`~torchcodec.transforms.DecoderTransform` and
             :class:`~torchvision.transforms.v2.Transform`
             objects. Read more about this parameter in :ref:`sphx_glr_generated_examples_decoding_transforms.py`.
+        output_dtype: TODO_HDR: add docs for output_dtype
         custom_frame_mappings (str, bytes, or file-like object, optional):
             Mapping of frames to their metadata, typically generated via ffprobe.
             This enables accurate frame seeking without requiring a full video scan.
@@ -168,6 +169,8 @@ def __init__(
         device: str | torch_device | None = None,
         seek_mode: Literal["exact", "approximate"] = "exact",
         transforms: Sequence[DecoderTransform | nn.Module] | None = None,
+        # TODO_HDR
+        output_dtype: torch.dtype | Literal["auto"] = torch.uint8,
         custom_frame_mappings: (
             str | bytes | io.RawIOBase | io.BufferedReader | None
         ) = None,
@@ -205,6 +208,15 @@ def __init__(
         if num_ffmpeg_threads is None:
             raise ValueError(f"{num_ffmpeg_threads = } should be an int.")
 
+        _DTYPE_TO_STR = {torch.uint8: "uint8", torch.float32: "float32"}
+        if output_dtype != "auto":
+            if output_dtype not in _DTYPE_TO_STR:
+                raise ValueError(
+                    f"Invalid output_dtype ({output_dtype}). "
+                    f"Supported values are torch.uint8, torch.float32, and 'auto'."
+                )
+            output_dtype = _DTYPE_TO_STR[output_dtype]
+
         device_variant = _get_cuda_backend()
         if device is None:
             device = str(torch.get_default_device())
@@ -224,6 +236,7 @@ def __init__(
             device_variant=device_variant,
             transforms=transforms,
             custom_frame_mappings=custom_frame_mappings_data,
+            output_dtype=output_dtype,
         )
 
         assert self.metadata.begin_stream_seconds is not None  # mypy.

test/generate_reference_resources.py

@@ -151,7 +151,7 @@ def generate_nasa_13013_references_by_index():
     # Note: The naming scheme used here must match the naming scheme used to load
     # tensors in ./utils.py.
     streams = [0, 3]
-    frames = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 20, 25, 30, 35, 386, 387, 388, 389]
+    frames = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 30, 35, 386, 387, 388, 389]
     for stream in streams:
         for frame in frames:
             generate_frame_by_index(NASA_VIDEO, frame_index=frame, stream_index=stream)
@@ -216,9 +216,18 @@ def generate_av1_video_references():
 
 def generate_hdr_references_rgb48():
     frames = [0, 5, 10]
-    for video in (NASA_VIDEO_HDR, TEST_SRC_2_720P_HDR, TEST_SRC_2_12BIT_HDR):
+    for video in (
+        NASA_VIDEO,
+        NASA_VIDEO_HDR,
+        TEST_SRC_2_720P_HDR,
+        TEST_SRC_2_12BIT_HDR,
+    ):
         for frame in frames:
-            generate_frame_by_index_rgb48(video, frame_index=frame, stream_index=0)
+            generate_frame_by_index_rgb48(
+                video,
+                frame_index=frame,
+                stream_index=video.default_stream_index,
+            )
 
 
 def generate_hdr_references_rgb24():

test/resources/nasa_13013.mp4.stream3.frame000000.rgb48.pt

@@ -42,11 +42,13 @@
     H265_10BITS,
     H265_VIDEO,
     in_fbcode,
+    IS_WINDOWS,
     make_video_decoder,
     NASA_AUDIO,
     NASA_AUDIO_MP3,
     NASA_AUDIO_MP3_44100,
     NASA_VIDEO,
+    NASA_VIDEO_HDR,
     NASA_VIDEO_ROTATED,
     needs_cuda,
     needs_ffmpeg_cli,
@@ -57,8 +59,10 @@
     SINE_MONO_S32_44100,
     SINE_MONO_S32_8000,
     TEST_NON_ZERO_START,
+    TEST_SRC_2_12BIT_HDR,
     TEST_SRC_2_720P,
     TEST_SRC_2_720P_H265,
+    TEST_SRC_2_720P_HDR,
     TEST_SRC_2_720P_MPEG4,
     TEST_SRC_2_720P_VP8,
     TEST_SRC_2_720P_VP9,
@@ -2360,6 +2364,145 @@ def test_beta_backend_still_supported_for_bc(self):
         dec[0]
         assert dec.cpu_fallback._backend == "CUDA"
 
+    @staticmethod
+    def _assert_float32_frame_matches_rgb48_ref(frame_data, asset, frame_index):
+        frame_as_uint16 = (frame_data * 65535).round().to(torch.uint16)
+        ref = asset.get_frame_data_by_index_rgb48(frame_index)
+        torch.testing.assert_close(frame_as_uint16, ref, rtol=0, atol=0)
+
+    @pytest.mark.xfail(
+        IS_WINDOWS and ffmpeg_major_version < 5,
+        reason="swscale YUV->RGB differs on Windows + FFmpeg 4",
+    )
+    @pytest.mark.parametrize(
+        "asset",
+        (NASA_VIDEO, NASA_VIDEO_HDR, TEST_SRC_2_720P_HDR, TEST_SRC_2_12BIT_HDR),
+    )
+    def test_output_dtype_uint8(self, asset):
+        decoder = VideoDecoder(asset.path, output_dtype=torch.uint8)
+        frame_indices = [0, 5, 10]
+        for frame_index in frame_indices:
+            frame = decoder[frame_index]
+            assert frame.dtype == torch.uint8
+            assert_frames_equal(frame.data, asset.get_frame_data_by_index(frame_index))
+
+    @pytest.mark.xfail(
+        IS_WINDOWS and ffmpeg_major_version < 5,
+        reason="swscale YUV->RGB48 differs on Windows + FFmpeg 4",
+    )
+    @pytest.mark.parametrize(
+        "asset",
+        (NASA_VIDEO, NASA_VIDEO_HDR, TEST_SRC_2_720P_HDR, TEST_SRC_2_12BIT_HDR),
+    )
+    def test_output_dtype_float32(self, asset):
+        decoder = VideoDecoder(asset.path, output_dtype=torch.float32)
+        frame_indices = [0, 5, 10]
+        for frame_index in frame_indices:
+            frame = decoder[frame_index]
+            assert frame.dtype == torch.float32
+
+            self._assert_float32_frame_matches_rgb48_ref(frame.data, asset, frame_index)
+
+    @pytest.mark.xfail(
+        IS_WINDOWS and ffmpeg_major_version < 5,
+        reason="swscale YUV->RGB48 differs on Windows + FFmpeg 4",
+    )
+    @pytest.mark.parametrize(
+        "asset, is_hdr",
+        (
+            (NASA_VIDEO, False),
+            (NASA_VIDEO_HDR, True),
+            (TEST_SRC_2_720P_HDR, True),
+            (TEST_SRC_2_12BIT_HDR, True),
+        ),
+    )
+    def test_output_dtype_auto(self, asset, is_hdr):
+        decoder = VideoDecoder(asset.path, output_dtype="auto")
+        frame_indices = [0, 5, 10]
+        for frame_index in frame_indices:
+            frame = decoder[frame_index]
+
+            if is_hdr:
+                assert frame.dtype == torch.float32
+                frame_as_uint16 = (frame.data * 65535).round().to(torch.uint16)
+                ref = asset.get_frame_data_by_index_rgb48(frame_index)
+                torch.testing.assert_close(frame_as_uint16, ref, rtol=0, atol=0)
+            else:
+                assert frame.dtype == torch.uint8
+                ref = asset.get_frame_data_by_index(frame_index)
+                assert_frames_equal(frame.data, ref)
+
+    @pytest.mark.xfail(
+        IS_WINDOWS and ffmpeg_major_version < 5,
+        reason="swscale YUV->RGB48 differs on Windows + FFmpeg 4",
+    )
+    @pytest.mark.parametrize(
+        "asset",
+        (NASA_VIDEO, NASA_VIDEO_HDR, TEST_SRC_2_720P_HDR, TEST_SRC_2_12BIT_HDR),
+    )
+    def test_output_dtype_float32_batch_apis(self, asset):
+        decoder = VideoDecoder(asset.path, output_dtype=torch.float32)
+        indices = [0, 5, 10]
+
+        # get_frame_at
+        self._assert_float32_frame_matches_rgb48_ref(
+            decoder.get_frame_at(0).data, asset, 0
+        )
+
+        # get_frames_at
+        frames = decoder.get_frames_at(indices)
+        for i, idx in enumerate(indices):
+            self._assert_float32_frame_matches_rgb48_ref(frames.data[i], asset, idx)
+
+        # get_frames_in_range
+        frames_range = decoder.get_frames_in_range(start=5, stop=11)
+        self._assert_float32_frame_matches_rgb48_ref(frames_range.data[0], asset, 5)
+        self._assert_float32_frame_matches_rgb48_ref(frames_range.data[5], asset, 10)
+
+    @pytest.mark.xfail(
+        IS_WINDOWS and ffmpeg_major_version < 5,
+        reason="swscale YUV->RGB48 differs on Windows + FFmpeg 4",
+    )
+    @pytest.mark.parametrize(
+        "asset",
+        (NASA_VIDEO, NASA_VIDEO_HDR, TEST_SRC_2_720P_HDR, TEST_SRC_2_12BIT_HDR),
+    )
+    def test_output_dtype_float32_pts_apis(self, asset):
+        decoder = VideoDecoder(asset.path, output_dtype=torch.float32)
+        indices = [0, 5, 10]
+
+        pts_seconds_ref = [decoder.get_frame_at(i).pts_seconds for i in indices]
+
+        # get_frame_played_at
+        for pts, idx in zip(pts_seconds_ref, indices):
+            frame = decoder.get_frame_played_at(pts)
+            self._assert_float32_frame_matches_rgb48_ref(frame.data, asset, idx)
+
+        # get_frames_played_in_range (full range)
+        frames = decoder.get_frames_played_in_range(
+            start_seconds=0,
+            stop_seconds=pts_seconds_ref[-1] + 1e-4,
+        )
+        for idx in indices:
+            self._assert_float32_frame_matches_rgb48_ref(frames.data[idx], asset, idx)
+
+        # get_frames_played_in_range (single-frame ranges)
+        for pts, idx in zip(pts_seconds_ref, indices):
+            frames = decoder.get_frames_played_in_range(
+                start_seconds=pts, stop_seconds=pts + 1e-4
+            )
+            self._assert_float32_frame_matches_rgb48_ref(frames.data[0], asset, idx)
+
+        # get_frames_played_at
+        frames = decoder.get_frames_played_at(pts_seconds_ref)
+        for i, idx in enumerate(indices):
+            self._assert_float32_frame_matches_rgb48_ref(frames.data[i], asset, idx)
+
+    @pytest.mark.parametrize("bad_dtype", (torch.float64, torch.int32, "not_a_dtype"))
+    def test_output_dtype_invalid(self, bad_dtype):
+        with pytest.raises(ValueError, match="Invalid output_dtype"):
+            VideoDecoder(NASA_VIDEO.path, output_dtype=bad_dtype)
+
 
 class TestAudioDecoder:
     @pytest.mark.parametrize(