diff --git a/src/torchcodec/_core/CudaDeviceInterface.cpp b/src/torchcodec/_core/CudaDeviceInterface.cpp index 8086d0b4b..62d8061aa 100644 --- a/src/torchcodec/_core/CudaDeviceInterface.cpp +++ b/src/torchcodec/_core/CudaDeviceInterface.cpp @@ -230,12 +230,23 @@ void CudaDeviceInterface::convertAVFrameToFrameOutput( auto start = std::chrono::high_resolution_clock::now(); NppStatus status; if (avFrame->colorspace == AVColorSpace::AVCOL_SPC_BT709) { - status = nppiNV12ToRGB_709CSC_8u_P2C3R( - input, - avFrame->linesize[0], - static_cast(dst.data_ptr()), - dst.stride(0), - oSizeROI); + if (avFrame->color_range == AVColorRange::AVCOL_RANGE_JPEG) { + // BT.709 full range + status = nppiNV12ToRGB_709HDTV_8u_P2C3R( + input, + avFrame->linesize[0], + static_cast(dst.data_ptr()), + dst.stride(0), + oSizeROI); + } else { + // BT.709 studio range + status = nppiNV12ToRGB_709CSC_8u_P2C3R( + input, + avFrame->linesize[0], + static_cast(dst.data_ptr()), + dst.stride(0), + oSizeROI); + } } else { status = nppiNV12ToRGB_8u_P2C3R( input, diff --git a/test/resources/full_range_709.mp4 b/test/resources/full_range_709.mp4 new file mode 100644 index 000000000..004028190 Binary files /dev/null and b/test/resources/full_range_709.mp4 differ diff --git a/test/test_decoders.py b/test/test_decoders.py index dcf9a1585..a4f11b990 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -25,6 +25,7 @@ assert_frames_equal, AV1_VIDEO, cpu_and_cuda, + FULL_COLOR_RANGE, get_ffmpeg_major_version, H265_VIDEO, in_fbcode, @@ -32,6 +33,7 @@ NASA_AUDIO_MP3, NASA_AUDIO_MP3_44100, NASA_VIDEO, + needs_cuda, SINE_MONO_S16, SINE_MONO_S32, SINE_MONO_S32_44100, @@ -1138,6 +1140,14 @@ def test_pts_to_dts_fallback(self, seek_mode): with pytest.raises(AssertionError, match="not equal"): torch.testing.assert_close(decoder[0], decoder[10]) + @needs_cuda + def test_full_range_bt709_video(self): + decoder_gpu = VideoDecoder(FULL_COLOR_RANGE.path, device="cuda") + decoder_cpu = VideoDecoder(FULL_COLOR_RANGE.path, device="cpu") + + assert_frames_equal(decoder_gpu[0].data.cpu(), decoder_cpu[0].data) + assert_frames_equal(decoder_gpu[10].data.cpu(), decoder_cpu[10].data) + class TestAudioDecoder: @pytest.mark.parametrize("asset", (NASA_AUDIO, NASA_AUDIO_MP3, SINE_MONO_S32)) diff --git a/test/utils.py b/test/utils.py index e3368e3f4..20675f01c 100644 --- a/test/utils.py +++ b/test/utils.py @@ -563,3 +563,12 @@ def sample_format(self) -> str: }, }, ) + +FULL_COLOR_RANGE = TestVideo( + filename="full_range_709.mp4", + default_stream_index=0, + stream_infos={ + 0: TestVideoStreamInfo(width=1280, height=720, num_color_channels=3), + }, + frames={0: {}}, # Not needed for now +)