diff --git a/packages/on_demand_video_decoder/data/pix_fmt_variants/hevc_hvc1_yuv444p.mp4 b/packages/on_demand_video_decoder/data/pix_fmt_variants/hevc_hvc1_yuv444p.mp4 new file mode 100644 index 00000000..3fb06c1c Binary files /dev/null and b/packages/on_demand_video_decoder/data/pix_fmt_variants/hevc_hvc1_yuv444p.mp4 differ diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncGopDecoder.cpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncGopDecoder.cpp index a5a80a36..35ff6e73 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncGopDecoder.cpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvBatchAsyncGopDecoder.cpp @@ -280,11 +280,11 @@ void PyNvBatchAsyncGopDecoder::build_yuv_frame(Pixel_Format fmt, size_t H, size_ break; // TODO(YUV444_16Bit): same as P016 — LoadDLPack rejects "|u2". case Pixel_Format_YUV444_16Bit: - out.views.push_back(CAIMemoryView{{H, W, 1}, {W, 1, 1}, "|u2", stream_id, dst_ptr, false}); + out.views.push_back(CAIMemoryView{{H, W, 1}, {W * 2, 2, 2}, "|u2", stream_id, dst_ptr, false}); out.views.push_back( - CAIMemoryView{{H, W, 1}, {W, 1, 1}, "|u2", stream_id, dst_ptr + 2 * H * W, false}); + CAIMemoryView{{H, W, 1}, {W * 2, 2, 2}, "|u2", stream_id, dst_ptr + 2 * H * W, false}); out.views.push_back( - CAIMemoryView{{H, W, 1}, {W, 1, 1}, "|u2", stream_id, dst_ptr + 4 * H * W, false}); + CAIMemoryView{{H, W, 1}, {W * 2, 2, 2}, "|u2", stream_id, dst_ptr + 4 * H * W, false}); break; default: // Only NV12 is currently supported. Returning a DecodedFrameExt with an empty extBuf diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp index f14e4444..109bc037 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp @@ -170,17 +170,23 @@ int PyNvGopDecoder::GetYUVFromFrame(NvDecoder* decoder, const uint8_t* pFrame, u (CUdeviceptr)(pFrame_buffer + width * height), false}); // todo: data+width*height assumes both planes are // contiguous. Actual NVENC allocation can have padding? + decoded_frame.views.push_back(CAIMemoryView{{height, width, 1}, + {width, 1, 1}, + "|u1", + reinterpret_cast(decoder->GetStream()), + (CUdeviceptr)(pFrame_buffer + 2 * width * height), + false}); } break; case Pixel_Format_YUV444_16Bit: { decoded_frame.views.push_back(CAIMemoryView{{height, width, 1}, - {width, 1, 1}, + {width * 2, 2, 2}, "|u2", reinterpret_cast(decoder->GetStream()), (CUdeviceptr)(pFrame_buffer), false}); decoded_frame.views.push_back( CAIMemoryView{{height, width, 1}, - {width, 1, 1}, + {width * 2, 2, 2}, "|u2", reinterpret_cast(decoder->GetStream()), (CUdeviceptr)(pFrame_buffer + 2 * (width * height)), @@ -188,7 +194,7 @@ int PyNvGopDecoder::GetYUVFromFrame(NvDecoder* decoder, const uint8_t* pFrame, u // contiguous. Actual NVENC allocation can have padding? decoded_frame.views.push_back( CAIMemoryView{{height, width, 1}, - {width, 1, 1}, + {width * 2, 2, 2}, "|u2", reinterpret_cast(decoder->GetStream()), (CUdeviceptr)(pFrame_buffer + 4 * (width * height)), diff --git a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvVideoReader.cpp b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvVideoReader.cpp index 9961e651..f216e9fc 100644 --- a/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvVideoReader.cpp +++ b/packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvVideoReader.cpp @@ -748,29 +748,37 @@ DecodedFrameExt PyNvVideoReader::returnYUVFrame(void* pFrame_buffer, void* pFram (CUdeviceptr)(pFrame_buffer + width * height), false}); // todo: data+width*height assumes both planes are // contiguous. Actual NVENC allocation can have padding? - } - case Pixel_Format_YUV444_16Bit: { frame.views.push_back(CAIMemoryView{{height, width, 1}, {width, 1, 1}, + "|u1", + reinterpret_cast(this->decoder->GetStream()), + (CUdeviceptr)(pFrame_buffer + 2 * width * height), + false}); + } break; + case Pixel_Format_YUV444_16Bit: { + frame.views.push_back(CAIMemoryView{{height, width, 1}, + {width * 2, 2, 2}, "|u2", reinterpret_cast(this->decoder->GetStream()), (CUdeviceptr)(pFrame_buffer), false}); frame.views.push_back(CAIMemoryView{{height, width, 1}, - {width, 1, 1}, + {width * 2, 2, 2}, "|u2", reinterpret_cast(this->decoder->GetStream()), (CUdeviceptr)(pFrame_buffer + 2 * (width * height)), false}); // todo: data+width*height assumes both planes are // contiguous. Actual NVENC allocation can have padding? frame.views.push_back(CAIMemoryView{{height, width, 1}, - {width, 1, 1}, + {width * 2, 2, 2}, "|u2", reinterpret_cast(this->decoder->GetStream()), (CUdeviceptr)(pFrame_buffer + 4 * (width * height)), false}); // todo: data+width*height assumes both planes are // contiguous. Actual NVENC allocation can have padding? - } + } break; + default: + throw std::runtime_error("[ERROR] Unsupported pixel format for YUV output"); } CUDA_DRVAPI_CALL(cuStreamSynchronize(this->decoder->GetStream())); return frame; diff --git a/packages/on_demand_video_decoder/tests/common/cases.py b/packages/on_demand_video_decoder/tests/common/cases.py index 20f6bdcc..ce052312 100644 --- a/packages/on_demand_video_decoder/tests/common/cases.py +++ b/packages/on_demand_video_decoder/tests/common/cases.py @@ -86,6 +86,15 @@ def resource_cases() -> Tuple[ResourceDecodeCase, ...]: torch.uint16, ((256, 256), (128, 128, 2)), ), + ResourceDecodeCase( + "hevc_hvc1_yuv444", + str(PIX_FMT_VARIANTS_DIR / "hevc_hvc1_yuv444p.mp4"), + 33, + "yuv", + "YUV444", + torch.uint8, + ((256, 256), (256, 256), (256, 256)), + ), ResourceDecodeCase( "vfr_h264", str(TEMPORAL_VARIANTS_DIR / "vfr_h264_yuv420p.mp4"), diff --git a/packages/on_demand_video_decoder/tests/test_pix_fmt_detection.py b/packages/on_demand_video_decoder/tests/test_pix_fmt_detection.py index 2c44a63e..0cdf073b 100644 --- a/packages/on_demand_video_decoder/tests/test_pix_fmt_detection.py +++ b/packages/on_demand_video_decoder/tests/test_pix_fmt_detection.py @@ -37,6 +37,11 @@ VARIANTS_DIR = os.path.join(utils.get_data_dir(), "pix_fmt_variants") +PIXEL_FORMAT_NV12 = 3 +PIXEL_FORMAT_YUV444 = 4 +PIXEL_FORMAT_P016 = 5 +PIXEL_FORMAT_YUV444_16BIT = 6 + # Each variant lists the filename, the codec_tag carried by the container, the # bit-depth the stream actually encodes, and the dtype that should appear on the @@ -48,6 +53,7 @@ ("hevc_hev1_yuv420p10le.mp4", "hev1", 10, 5, "|u2", ((256, 256, 1), (128, 128, 2))), ("hevc_hvc1_yuv420p.mp4", "hvc1", 8, 3, "|u1", ((256, 256, 1), (128, 128, 2))), ("hevc_hvc1_yuv420p10le.mp4", "hvc1", 10, 5, "|u2", ((256, 256, 1), (128, 128, 2))), + ("hevc_hvc1_yuv444p.mp4", "hvc1", 8, 4, "|u1", ((256, 256, 1),) * 3), ("h264_avc1_yuv420p.mp4", "avc1", 8, 3, "|u1", ((256, 256, 1), (128, 128, 2))), ] @@ -59,6 +65,34 @@ def _video_path(name): return path +def _expected_plane_strides(pixel_format, luma_width, bytes_per_sample): + if pixel_format in (PIXEL_FORMAT_NV12, PIXEL_FORMAT_P016): + # Semi-planar 4:2:0: one Y plane followed by one interleaved UV plane. + return ( + ( + luma_width * bytes_per_sample, + bytes_per_sample, + bytes_per_sample, + ), + ( + luma_width * bytes_per_sample, + 2 * bytes_per_sample, + bytes_per_sample, + ), + ) + + if pixel_format in (PIXEL_FORMAT_YUV444, PIXEL_FORMAT_YUV444_16BIT): + # Planar 4:4:4: Y, U, and V are three full-resolution planes. + plane_strides = ( + luma_width * bytes_per_sample, + bytes_per_sample, + bytes_per_sample, + ) + return (plane_strides,) * 3 + + raise ValueError(f"unsupported pixel format in stride check: {pixel_format}") + + @pytest.mark.parametrize( "filename, codec_tag, bit_depth, expected_format, expected_dtype, expected_shapes", VARIANTS, @@ -97,17 +131,10 @@ def test_decode_from_gop_round_trip( expected_bytes_per_sample = 2 if bit_depth >= 10 else 1 luma_width = expected_shapes[0][1] - expected_strides = ( - ( - luma_width * expected_bytes_per_sample, - expected_bytes_per_sample, - expected_bytes_per_sample, - ), - ( - luma_width * expected_bytes_per_sample, - 2 * expected_bytes_per_sample, - expected_bytes_per_sample, - ), + expected_strides = _expected_plane_strides( + expected_format, + luma_width, + expected_bytes_per_sample, ) actual_strides = tuple(tuple(plane.__cuda_array_interface__["strides"]) for plane in planes) assert (