diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index d70c8bf66..c2f7f9076 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -4,6 +4,7 @@ // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. +#include #include #include #include @@ -343,7 +344,7 @@ std::optional get_nvdec_surface_format( void standalone_frame_free_callback( [[maybe_unused]] void* opaque, uint8_t* data) { - delete reinterpret_cast(data); + delete reinterpret_cast(data); } class CudaContextGuard { @@ -964,7 +965,12 @@ UniqueAVFrame BetaCudaDeviceInterface::convert_cuda_frame_to_av_frame( av_frame->data[1] = plane(1); av_frame->data[2] = is_444 ? plane(2) : nullptr; av_frame->data[3] = nullptr; - // TODO_API_BREAKDOWN CC P2: Check range before cast? + STD_TORCH_CHECK( + pitch <= static_cast(std::numeric_limits::max()), + "NVDEC returned a pitch of ", + pitch, + " bytes, which doesn't fit in an AVFrame line size. This should never " + "happen, please report."); av_frame->linesize[0] = static_cast(pitch); av_frame->linesize[1] = static_cast(pitch); av_frame->linesize[2] = is_444 ? static_cast(pitch) : 0; @@ -1000,12 +1006,12 @@ void BetaCudaDeviceInterface::make_frame_standalone(UniqueAVFrame& av_frame) { storage = copy_nvdec_surface(av_frame, current_stream); } - auto attached_data = new StandAloneFrameAttachedData(); + auto attached_data = new OwnedFrameStorage(); attached_data->frame_ready.record(current_stream); attached_data->storage = std::move(storage); av_frame->opaque_ref = av_buffer_create( reinterpret_cast(attached_data), - sizeof(StandAloneFrameAttachedData), + sizeof(OwnedFrameStorage), standalone_frame_free_callback, nullptr, 0); @@ -1065,8 +1071,7 @@ std::optional BetaCudaDeviceInterface::get_frame_storage( // for those users who would like to consume the frame with their own // consumer, i.e. not using the ColorConverter: they need to call // frame.storage.record_stream(color_conversion_stream) themselves. - return reinterpret_cast( - av_frame.opaque_ref->data) + return reinterpret_cast(av_frame.opaque_ref->data) ->storage; } @@ -1233,10 +1238,21 @@ GpuFrameAndStorage BetaCudaDeviceInterface::upload_cpu_frame_to_gpu( cpu_frame.colorspace, width, height, - target_pix_fmt); + target_pix_fmt, + // The frame keeps its range tag through the upload, so the samples must + // keep the range that tag names. Left to itself, swscale writes limited + // range into a YUV target whatever it read, and the color conversion + // would then expand a full-range source a second time. + cpu_frame.color_range); if (!sws_context_ || prev_sws_config_ != sws_config) { - sws_context_ = create_sws_context(sws_config, SWS_BILINEAR); + // Nothing is rescaled here, so the flags only pick how chroma is + // resampled, which happens when the source is subsampled more finely than + // the target surface (4:2:2 into 4:4:4, say). SWS_POINT replicates it, + // which is what the CPU converter does on its way to RGB - interpolating + // instead would invent chroma the CPU never sees, and show up as colored + // fringes along sharp edges. + sws_context_ = create_sws_context(sws_config, SWS_POINT); prev_sws_config_ = sws_config; } @@ -1323,6 +1339,13 @@ GpuFrameAndStorage BetaCudaDeviceInterface::upload_cpu_frame_to_gpu( "Failed to copy frame properties: ", get_ffmpeg_error_string_from_error_code(ret)); + // AVCOL_SPC_RGB says "these planes are RGB", which the planes we just wrote + // aren't. Name the matrix swscale encoded them with instead: it maps + // AVCOL_SPC_RGB, like any colorspace it doesn't know, to its BT.601 default. + if (cpu_frame.colorspace == AVCOL_SPC_RGB) { + gpu_frame->colorspace = AVCOL_SPC_SMPTE170M; + } + return {std::move(gpu_frame), std::move(storage)}; } @@ -1369,8 +1392,8 @@ void BetaCudaDeviceInterface::convert_av_frame_to_frame_output( gpu_frame.opaque_ref != nullptr, "ColorConverter received a non-standalone frame; frames fed to a " "standalone ColorConverter must come from a PacketDecoder."); - auto attached_data = reinterpret_cast( - gpu_frame.opaque_ref->data); + auto attached_data = + reinterpret_cast(gpu_frame.opaque_ref->data); attached_data->frame_ready.make_stream_wait(current_stream); } else { STD_TORCH_CHECK( diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index ede4068ef..fc2b82cbd 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.h @@ -34,9 +34,8 @@ #include "nvcuvid_include/nvcuvid.h" namespace facebook::torchcodec { -// TODO_API_BREAKDOWN P2: the name says "standalone", but this is really about -// owning a GPU buffer. Find one that covers both. -struct StandAloneFrameAttachedData { +// The buffer a frame owns its samples in, hung off the AVFrame as opaque data. +struct OwnedFrameStorage { // Marks the point where the copy (or upload) that filled `storage` was // enqueued. A consumer on another stream must wait on it. CudaEvent frame_ready; diff --git a/src/torchcodec/_core/FFMPEGCommon.cpp b/src/torchcodec/_core/FFMPEGCommon.cpp index 03f5c8464..43ebb3cd4 100644 --- a/src/torchcodec/_core/FFMPEGCommon.cpp +++ b/src/torchcodec/_core/FFMPEGCommon.cpp @@ -904,14 +904,16 @@ SwsConfig::SwsConfig( AVColorSpace input_colorspace, int output_width, int output_height, - AVPixelFormat output_format) + AVPixelFormat output_format, + AVColorRange output_color_range) : input_width(input_width), input_height(input_height), input_format(input_format), input_colorspace(input_colorspace), output_width(output_width), output_height(output_height), - output_format(output_format) {} + output_format(output_format), + output_color_range(output_color_range) {} bool SwsConfig::operator==(const SwsConfig& other) const { return input_width == other.input_width && @@ -920,7 +922,8 @@ bool SwsConfig::operator==(const SwsConfig& other) const { input_colorspace == other.input_colorspace && output_width == other.output_width && output_height == other.output_height && - output_format == other.output_format; + output_format == other.output_format && + output_color_range == other.output_color_range; } bool SwsConfig::operator!=(const SwsConfig& other) const { @@ -957,6 +960,13 @@ UniqueSwsContext create_sws_context( &saturation); STD_TORCH_CHECK(ret != -1, "sws_getColorspaceDetails returned -1"); + // swscale spells a range as an int: 1 is full, 0 is limited. FFmpeg names + // those AVCOL_RANGE_JPEG and AVCOL_RANGE_MPEG, after the two worlds they come + // from - JPEG is the full one. + if (sws_config.output_color_range != AVCOL_RANGE_UNSPECIFIED) { + dst_range = sws_config.output_color_range == AVCOL_RANGE_JPEG; + } + const int* colorspace_table = sws_getCoefficients(sws_config.input_colorspace); ret = sws_setColorspaceDetails( diff --git a/src/torchcodec/_core/FFMPEGCommon.h b/src/torchcodec/_core/FFMPEGCommon.h index 12e63d2c0..53d68bbee 100644 --- a/src/torchcodec/_core/FFMPEGCommon.h +++ b/src/torchcodec/_core/FFMPEGCommon.h @@ -406,6 +406,10 @@ struct SwsConfig { int output_width = 0; int output_height = 0; AVPixelFormat output_format = AV_PIX_FMT_NONE; + // swscale derives the output range from the output pixel format alone: YUV + // gets limited range. AVCOL_RANGE_UNSPECIFIED keeps that; anything else + // overrides it. + AVColorRange output_color_range = AVCOL_RANGE_UNSPECIFIED; SwsConfig() = default; SwsConfig( @@ -415,7 +419,8 @@ struct SwsConfig { AVColorSpace input_colorspace, int output_width, int output_height, - AVPixelFormat output_format); + AVPixelFormat output_format, + AVColorRange output_color_range = AVCOL_RANGE_UNSPECIFIED); bool operator==(const SwsConfig& other) const; bool operator!=(const SwsConfig& other) const; diff --git a/test/resources/testsrc2_gbrp_hevc.mp4 b/test/resources/testsrc2_gbrp_hevc.mp4 new file mode 100644 index 000000000..f5b96cd99 Binary files /dev/null and b/test/resources/testsrc2_gbrp_hevc.mp4 differ diff --git a/test/resources/testsrc2_gray_hevc.mp4 b/test/resources/testsrc2_gray_hevc.mp4 new file mode 100644 index 000000000..20724cb5d Binary files /dev/null and b/test/resources/testsrc2_gray_hevc.mp4 differ diff --git a/test/resources/testsrc2_yuva420p_ffv1.mkv b/test/resources/testsrc2_yuva420p_ffv1.mkv new file mode 100644 index 000000000..4e521a1a0 Binary files /dev/null and b/test/resources/testsrc2_yuva420p_ffv1.mkv differ diff --git a/test/test_decoders.py b/test/test_decoders.py index 273ad8f51..8aa255cfc 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -144,6 +144,9 @@ TESTSRC2_444_12BIT_HEVC, TESTSRC2_444_8BIT_HEVC, TESTSRC2_AV1_10BIT, + TESTSRC2_FULL_RANGE_422, + TESTSRC2_GBRP_HEVC, + TESTSRC2_GRAY_HEVC, TESTSRC2_ODD_HEIGHT_444, TESTSRC2_ODD_HEIGHT_AND_WIDTH_444, TESTSRC2_ODD_HEIGHT_AND_WIDTH_444_10BIT, @@ -156,6 +159,7 @@ TESTSRC2_ODD_WIDTH_MPEG2, TESTSRC2_ODD_WIDTH_VP9, TESTSRC2_ODD_WIDTH_VP9_10BIT, + TESTSRC2_YUVA420P_FFV1, TRANSPARENT_GIF, UNSEEKABLE_SWF, WAV_ODD_DATA_TRAILING_CHUNK, @@ -2365,6 +2369,36 @@ def test_nvdec_cpu_fallback_yuv444(self, tmp_path): # our kernel, which truncates where swscale rounds. torch.testing.assert_close(cpu_frames, cuda_frames.cpu(), rtol=0, atol=1) + @needs_cuda + @pytest.mark.parametrize( + "video", + ( + TESTSRC2_GRAY_HEVC, + TESTSRC2_GBRP_HEVC, + TESTSRC2_YUVA420P_FFV1, + TESTSRC2_FULL_RANGE_422, + ), + ids=lambda video: video.path.stem, + ) + def test_cpu_fallback_matches_cpu(self, video): + # NVDEC decodes none of these, so CUDA decodes them on the CPU and + # uploads them in an NVDEC surface format. Monochrome, planar RGB, alpha + # and 4:2:2 all convert to something those formats hold, and all but the + # FFV1 one are full range, which is what a conversion that quietly + # narrowed them to limited range would squash. + num_frames = 5 + cpu_decoder = VideoDecoder(video.path, device="cpu") + cuda_decoder = VideoDecoder(video.path, device="cuda") + assert cuda_decoder.cpu_fallback + + cpu_frames = cpu_decoder[:num_frames] + cuda_frames = cuda_decoder[:num_frames].cpu() + + # A couple of levels for the color-conversion kernel, and a couple more + # for the planar RGB source, whose samples make a round trip through + # 8-bit YUV that the CPU never puts them through. + torch.testing.assert_close(cuda_frames, cpu_frames, atol=3, rtol=0) + @needs_cuda def test_nvdec_cuda_interface_error(self): with pytest.raises(RuntimeError, match="torch_parse_device_string"): @@ -3552,6 +3586,10 @@ class _PlanesCase(NamedTuple): bit_depth: int cpu_pix_fmt: str cuda_pix_fmt: str + # One per component of the pixel format, so three for YUV and RGB, one for + # grayscale, and one more when the format has an alpha component. + cpu_num_planes: int = 3 + cuda_num_planes: int = 3 # FFmpeg 6 added P012. Before that, NVDEC's 12-bit surface can only be # described as p016le, which claims 16 bits instead of 12. Set this for the # sources that hit it: same samples either way (they're msb-aligned, so @@ -3562,6 +3600,9 @@ class _PlanesCase(NamedTuple): def pix_fmt(self, device): return self.cuda_pix_fmt if device == "cuda" else self.cpu_pix_fmt + def num_planes(self, device): + return self.cuda_num_planes if device == "cuda" else self.cpu_num_planes + # Sources with more than 8 bits per sample. _HDR_VIDEOS = ( @@ -3577,7 +3618,8 @@ def pix_fmt(self, device): # Videos spanning the pixel-format axes RawFrame.planes has to handle: 4:2:0 vs # 4:4:4 chroma, even vs odd dims (chroma rounds up), and 8- vs 10-/12-bit -# (uint8 vs uint16 planes). All are YUV, so planes are (Y, U, V). +# (uint8 vs uint16 planes). All are YUV, so planes are (Y, U, V) - the sources +# whose frames aren't three YUV planes are in _NON_YUV_PLANES_VIDEOS below. _PLANES_VIDEOS = ( _PlanesCase(NASA_VIDEO, 8, "yuv420p", "nv12"), # even dims _PlanesCase(TESTSRC2_ODD_HEIGHT_AND_WIDTH_VP9, 8, "yuv420p", "nv12"), # odd @@ -3604,6 +3646,24 @@ def pix_fmt(self, device): ) +# Sources whose frames aren't three YUV planes on the CPU. NVDEC decodes none of +# them - monochrome, planar RGB and FFV1 all send it to the CPU fallback - and +# the fallback converts to an NVDEC surface format before uploading, so on CUDA +# they are three YUV planes like everything else. That conversion is what +# RawFrame.pix_fmt promises ("on CUDA it is always an NVDEC surface format"). It +# keeps the pixels where they were - see test_cpu_fallback_matches_cpu - but not +# what YUV has no room for: grayscale gains neutral chroma, and alpha is dropped +# outright. +_NON_YUV_PLANES_VIDEOS = ( + _PlanesCase(TESTSRC2_GRAY_HEVC, 8, "gray", "nv12", cpu_num_planes=1), + # Planar RGB. The planes come out (R, G, B), which is *not* the order the + # format stores them in: FFmpeg's gbrp is green, blue, red. + _PlanesCase(TESTSRC2_GBRP_HEVC, 8, "gbrp", "yuv444p"), + # Alpha, which is full size like luma rather than subsampled like chroma. + _PlanesCase(TESTSRC2_YUVA420P_FFV1, 8, "yuva420p", "nv12", cpu_num_planes=4), +) + + def _planes_ids(case): return case.video.path.stem @@ -4211,7 +4271,9 @@ def test_device_torch_device_instance(self, device): assert frame.planes[0].device.type == device assert converter.convert(frame).data.device.type == device - @pytest.mark.parametrize("case", _PLANES_VIDEOS, ids=_planes_ids) + @pytest.mark.parametrize( + "case", _PLANES_VIDEOS + _NON_YUV_PLANES_VIDEOS, ids=_planes_ids + ) @pytest.mark.parametrize("device", _block_devices()) def test_planes_structure(self, case, device): # planes shape/dtype/device and the accompanying metadata. @@ -4233,12 +4295,13 @@ def test_planes_structure(self, case, device): ) assert pix_fmt == expected_pix_fmt - assert frame.colorspace in ("bt709", "bt2020nc", "smpte170m", "unknown") + # "gbr" is what a planar RGB frame reports: its samples are already RGB, + # so there is no YUV matrix to name. + assert frame.colorspace in ("bt709", "bt2020nc", "smpte170m", "gbr", "unknown") assert frame.color_range in ("tv", "pc", "unknown") # FFmpeg has only these # All planes are 2D views living on the frame's own device. - # TODO_API_BREAKDOWN DESIGN P1: Can there be more planes? Should test? - assert len(planes) == 3 + assert len(planes) == case.num_planes(device) for plane in planes: assert plane.ndim == 2 assert plane.device.type == device @@ -4257,18 +4320,25 @@ def test_planes_structure(self, case, device): for plane in planes: assert (plane.to(torch.int32) & unused_low_bits).count_nonzero() == 0 - Y, U, V = planes height, width = converter.convert(frame).data.shape[1:] - assert Y.shape == (height, width) == (frame.height, frame.width) + # The first plane is always full size: luma, or red for a planar RGB + # format. So is a trailing alpha one, when the format has it. + assert planes[0].shape == (height, width) == (frame.height, frame.width) + if len(planes) == 4: + assert planes[3].shape == (height, width) # Below is just a fancy way to divide by 2 accounting for odd sizes, # matching the FFmpeg logic - log2_h, log2_w = (0, 0) if "444" in pix_fmt else (1, 1) + subsampled = not (pix_fmt.startswith("gbr") or "444" in pix_fmt) + log2_h, log2_w = (1, 1) if subsampled else (0, 0) expected_chroma_shape = ( (height + (1 << log2_h) - 1) >> log2_h, (width + (1 << log2_w) - 1) >> log2_w, ) - assert U.shape == V.shape == expected_chroma_shape + # Planes 1 and 2 are the subsampled ones for YUV, and full-size green + # and blue for planar RGB. Grayscale has neither. + for plane in planes[1:3]: + assert plane.shape == expected_chroma_shape @pytest.mark.parametrize("device", _block_devices()) def test_planes_are_not_rotated_but_color_conversion_rotates(self, device): @@ -4504,6 +4574,50 @@ def test_cpu_fallback_is_on_cuda(self, video, expected_pix_fmt): assert frame.pix_fmt == expected_pix_fmt assert all(plane.device.type == "cuda" for plane in frame.planes) + @pytest.mark.needs_cuda + @pytest.mark.parametrize( + "video, expected_colorspace, has_luma", + ( + pytest.param(TESTSRC2_GRAY_HEVC, "unknown", True, id="gray"), + # "gbr" describes RGB planes, which the uploaded frame doesn't have. + # It names the YUV matrix its planes were encoded with instead. + pytest.param(TESTSRC2_GBRP_HEVC, "smpte170m", False, id="gbrp"), + pytest.param(TESTSRC2_FULL_RANGE_422, "unknown", True, id="422"), + ), + ) + def test_cpu_fallback_upload_keeps_full_range( + self, video, expected_colorspace, has_luma + ): + # Full-range sources NVDEC can't decode, so they go through the CPU + # fallback and its conversion to an NVDEC surface format. The samples + # stay full range across that conversion: narrowing them to limited + # range while the frame still says "pc" would have the color conversion + # stretch them a second time. + cpu_frame, cpu_converter = self._first_frame(video.path, "cpu") + cuda_frame, cuda_converter = self._first_frame(video.path, "cuda") + + assert cpu_frame.color_range == "pc" + assert cuda_frame.color_range == "pc" + assert cuda_frame.colorspace == expected_colorspace + + if has_luma: + # A source that already has luma keeps it sample for sample: only + # its chroma is touched. + height, width = cpu_frame.planes[0].shape + torch.testing.assert_close( + cuda_frame.planes[0][:height, :width].cpu(), + cpu_frame.planes[0], + atol=0, + rtol=0, + ) + + torch.testing.assert_close( + cuda_converter.convert(cuda_frame).data.cpu(), + cpu_converter.convert(cpu_frame).data, + atol=3, + rtol=0, + ) + @pytest.mark.parametrize( "pix_fmt, codec, container", ( diff --git a/test/utils.py b/test/utils.py index b1db9eae2..ed2234851 100644 --- a/test/utils.py +++ b/test/utils.py @@ -1274,6 +1274,59 @@ def get_empty_chw_tensor(self, *, stream_index: int) -> torch.Tensor: frames={0: {}}, ) +# The sources whose frames don't come out as three YUV planes. libx264 accepts +# -pix_fmt gray but silently encodes 4:2:0 anyway, hence libx265 here. Even +# dimensions, because both encoders below round odd ones down. +# ffmpeg -f lavfi -i "testsrc2=size=320x240:rate=25:duration=1" \ +# -vf format=gray -c:v libx265 -tag:v hvc1 testsrc2_gray_hevc.mp4 +TESTSRC2_GRAY_HEVC = TestVideo( + filename="testsrc2_gray_hevc.mp4", + default_stream_index=0, + stream_infos={ + 0: TestVideoStreamInfo(width=320, height=240, num_color_channels=3), + }, + frames={0: {}}, +) + +# ffmpeg -f lavfi -i "testsrc2=size=321x241:rate=25:duration=1,format=rgb24" \ +# -vf format=gbrp -c:v libx265 -tag:v hvc1 testsrc2_gbrp_hevc.mp4 +TESTSRC2_GBRP_HEVC = TestVideo( + filename="testsrc2_gbrp_hevc.mp4", + default_stream_index=0, + stream_infos={ + 0: TestVideoStreamInfo(width=321, height=241, num_color_channels=3), + }, + frames={0: {}}, +) + +# FFV1 is lossless, so this one is a fifth of a second rather than a full one. +# VP9's alpha is not an option: it rides in a separate layer, and the frames the +# decoder produces are plain yuv420p. +# ffmpeg -f lavfi -i "testsrc2=size=320x240:rate=25:duration=0.2" \ +# -vf format=yuva420p -c:v ffv1 testsrc2_yuva420p_ffv1.mkv +TESTSRC2_YUVA420P_FFV1 = TestVideo( + filename="testsrc2_yuva420p_ffv1.mkv", + default_stream_index=0, + stream_infos={ + 0: TestVideoStreamInfo(width=320, height=240, num_color_channels=3), + }, + frames={0: {}}, +) + +# Full range (pc) 4:2:2, i.e. neither the range nor the chroma layout that NVDEC +# surfaces come in. yuvj422p is what -pix_fmt yuvj422p and -pix_fmt yuv422p +# -color_range pc both produce. +# ffmpeg -f lavfi -i "testsrc2=size=320x240:rate=25:duration=1" \ +# -c:v libx264 -pix_fmt yuvj422p testsrc2_full_range_422.mp4 +TESTSRC2_FULL_RANGE_422 = TestVideo( + filename="testsrc2_full_range_422.mp4", + default_stream_index=0, + stream_infos={ + 0: TestVideoStreamInfo(width=320, height=240, num_color_channels=3), + }, + frames={0: {}}, +) + # ffmpeg -f lavfi -i "testsrc2=size=321x240:rate=25:duration=1,format=rgb24" \ # -c:v libvpx-vp9 -pix_fmt yuv420p -b:v 1M testsrc2_odd_width_vp9.mp4 TESTSRC2_ODD_WIDTH_VP9 = TestVideo(