From a7e5dfc567063b1d56674a5127bc5a270a9fac9c Mon Sep 17 00:00:00 2001 From: nightcityblade Date: Wed, 13 May 2026 23:29:58 +0800 Subject: [PATCH 1/2] Avoid deprecated nvcv_image conversion Signed-off-by: nightcityblade --- nemo_curator/utils/nvcodec_utils.py | 3 +-- tests/utils/test_nvcodec_utils.py | 18 ++++-------------- 2 files changed, 5 insertions(+), 16 deletions(-) diff --git a/nemo_curator/utils/nvcodec_utils.py b/nemo_curator/utils/nvcodec_utils.py index fefb361634..c0e68e344a 100644 --- a/nemo_curator/utils/nvcodec_utils.py +++ b/nemo_curator/utils/nvcodec_utils.py @@ -268,8 +268,7 @@ def generate_decoded_frames(self) -> list[torch.Tensor]: for packet in self.nvDemux: list_frames = self.nvDec.Decode(packet) for decoded_frame in list_frames: - # TODO: Remove the use of nvcv_image. It's deprecated - cvcuda_tensor = cvcuda.as_tensor(cvcuda.as_image(decoded_frame.nvcv_image(), cvcuda.Format.U8)) + cvcuda_tensor = cvcuda.as_tensor(decoded_frame, "NCHW") if cvcuda_tensor.layout == "NCHW": nchw_shape = cvcuda_tensor.shape nhwc_shape = (nchw_shape[0], nchw_shape[2], nchw_shape[3], nchw_shape[1]) diff --git a/tests/utils/test_nvcodec_utils.py b/tests/utils/test_nvcodec_utils.py index 181a7a08ca..c381cf8980 100644 --- a/tests/utils/test_nvcodec_utils.py +++ b/tests/utils/test_nvcodec_utils.py @@ -667,16 +667,11 @@ def test_generate_decoded_frames_with_frames(self, mock_cvcuda: Any, mock_torch: # Mock packet and decoded frame mock_packet = Mock() mock_decoded_frame = Mock() - mock_decoded_frame.nvcv_image.return_value = Mock() - # Mock tensor operations mock_cvcuda_tensor = Mock() mock_cvcuda_tensor.layout = "NCHW" mock_cvcuda_tensor.shape = (1, 3, 480, 640) # NCHW format mock_cvcuda.as_tensor.return_value = mock_cvcuda_tensor - mock_cvcuda.as_image.return_value = Mock() - mock_cvcuda.Format.U8 = Mock() - # Mock torch tensor mock_torch_nhwc = Mock() mock_torch.empty.return_value = mock_torch_nhwc @@ -686,6 +681,8 @@ def test_generate_decoded_frames_with_frames(self, mock_cvcuda: Any, mock_torch: # Setup side effect for as_tensor to return different mocks def as_tensor_side_effect(*args: Any, **_kwargs: Any) -> Any: + if args[0] is mock_decoded_frame: + return mock_cvcuda_tensor if len(args) == 2 and isinstance(args[1], str): return mock_cvcuda_nhwc return mock_cvcuda_tensor @@ -735,15 +732,10 @@ def test_generate_decoded_frames_unexpected_layout( # Mock packet and decoded frame mock_packet = Mock() mock_decoded_frame = Mock() - mock_decoded_frame.nvcv_image.return_value = Mock() - # Mock tensor with unexpected layout mock_cvcuda_tensor = Mock() mock_cvcuda_tensor.layout = "NHWC" # Unexpected layout - should be NCHW mock_cvcuda.as_tensor.return_value = mock_cvcuda_tensor - mock_cvcuda.as_image.return_value = Mock() - mock_cvcuda.Format.U8 = Mock() - # Mock demux iteration mock_demux.__iter__ = Mock(return_value=iter([mock_packet])) mock_decoder.Decode.return_value = [mock_decoded_frame] @@ -780,7 +772,6 @@ def test_generate_decoded_frames_partial_batch(self, mock_cvcuda: Any, mock_torc mock_frames = [] for _ in range(2): # Only 2 frames, less than batch_size of 4 frame = Mock() - frame.nvcv_image.return_value = Mock() mock_frames.append(frame) # Mock tensor operations @@ -788,15 +779,14 @@ def test_generate_decoded_frames_partial_batch(self, mock_cvcuda: Any, mock_torc mock_cvcuda_tensor.layout = "NCHW" mock_cvcuda_tensor.shape = (1, 3, 480, 640) mock_cvcuda.as_tensor.return_value = mock_cvcuda_tensor - mock_cvcuda.as_image.return_value = Mock() - mock_cvcuda.Format.U8 = Mock() - mock_torch_nhwc = Mock() mock_torch.empty.return_value = mock_torch_nhwc mock_cvcuda_nhwc = Mock() # Setup side effect for as_tensor to return different mocks def as_tensor_side_effect(*args: Any, **_kwargs: Any) -> Any: + if any(args[0] is frame for frame in mock_frames): + return mock_cvcuda_tensor if len(args) == 2 and isinstance(args[1], str): return mock_cvcuda_nhwc return mock_cvcuda_tensor From a86ac7316ee1b5e171ccdb843bc7de18b406a1a6 Mon Sep 17 00:00:00 2001 From: nightcityblade Date: Thu, 14 May 2026 23:08:20 +0800 Subject: [PATCH 2/2] Remove redundant decoded frame layout guard --- nemo_curator/utils/nvcodec_utils.py | 26 ++++++++----------- tests/utils/test_nvcodec_utils.py | 40 ----------------------------- 2 files changed, 11 insertions(+), 55 deletions(-) diff --git a/nemo_curator/utils/nvcodec_utils.py b/nemo_curator/utils/nvcodec_utils.py index c0e68e344a..082b7e91d6 100644 --- a/nemo_curator/utils/nvcodec_utils.py +++ b/nemo_curator/utils/nvcodec_utils.py @@ -269,21 +269,17 @@ def generate_decoded_frames(self) -> list[torch.Tensor]: list_frames = self.nvDec.Decode(packet) for decoded_frame in list_frames: cvcuda_tensor = cvcuda.as_tensor(decoded_frame, "NCHW") - if cvcuda_tensor.layout == "NCHW": - nchw_shape = cvcuda_tensor.shape - nhwc_shape = (nchw_shape[0], nchw_shape[2], nchw_shape[3], nchw_shape[1]) - torch_nhwc = torch.empty( - nhwc_shape, - dtype=torch.uint8, - device=f"cuda:{self.device_id}", - ) - cvcuda_nhwc = cvcuda.as_tensor(torch_nhwc.cuda(self.device_id), "NHWC") - cvcuda.reformat_into(cvcuda_nhwc, cvcuda_tensor, stream=self.cvcuda_stream) - # Push the decoded frame with the reformatted frame to keep it alive. - self.input_frame_list.put(torch_nhwc) - else: - error_msg = "Unexpected tensor layout, NCHW expected." - raise ValueError(error_msg) + nchw_shape = cvcuda_tensor.shape + nhwc_shape = (nchw_shape[0], nchw_shape[2], nchw_shape[3], nchw_shape[1]) + torch_nhwc = torch.empty( + nhwc_shape, + dtype=torch.uint8, + device=f"cuda:{self.device_id}", + ) + cvcuda_nhwc = cvcuda.as_tensor(torch_nhwc.cuda(self.device_id), "NHWC") + cvcuda.reformat_into(cvcuda_nhwc, cvcuda_tensor, stream=self.cvcuda_stream) + # Push the decoded frame with the reformatted frame to keep it alive. + self.input_frame_list.put(torch_nhwc) self.local_frame_index = self.local_frame_index + 1 self.decoded_frame_cnt = self.decoded_frame_cnt + 1 diff --git a/tests/utils/test_nvcodec_utils.py b/tests/utils/test_nvcodec_utils.py index c381cf8980..431ace3dd4 100644 --- a/tests/utils/test_nvcodec_utils.py +++ b/tests/utils/test_nvcodec_utils.py @@ -711,46 +711,6 @@ def as_tensor_side_effect(*args: Any, **_kwargs: Any) -> Any: # Should return the processed frames assert result == [mock_torch_nhwc] - @patch("nemo_curator.utils.nvcodec_utils.Nvc") - @patch("nemo_curator.utils.nvcodec_utils.torch") - @patch("nemo_curator.utils.nvcodec_utils.cvcuda") - def test_generate_decoded_frames_unexpected_layout( - self, mock_cvcuda: Any, _mock_torch: Any, mock_nvc: Any - ) -> None: - """Test generate_decoded_frames with unexpected tensor layout.""" - # Setup mocks - mock_demux = Mock() - mock_demux.Width.return_value = 640 - mock_demux.Height.return_value = 480 - mock_demux.GetNvCodecId.return_value = "H264" - mock_nvc.PyNvDemuxer.return_value = mock_demux - - mock_decoder = Mock() - mock_decoder.GetPixelFormat.return_value = "NV12" - mock_nvc.CreateDecoder.return_value = mock_decoder - - # Mock packet and decoded frame - mock_packet = Mock() - mock_decoded_frame = Mock() - # Mock tensor with unexpected layout - mock_cvcuda_tensor = Mock() - mock_cvcuda_tensor.layout = "NHWC" # Unexpected layout - should be NCHW - mock_cvcuda.as_tensor.return_value = mock_cvcuda_tensor - # Mock demux iteration - mock_demux.__iter__ = Mock(return_value=iter([mock_packet])) - mock_decoder.Decode.return_value = [mock_decoded_frame] - - decoder = NvVideoDecoder( - enc_file="test_video.mp4", - device_id=self.device_id, - batch_size=self.batch_size, - cuda_ctx=self.mock_cuda_ctx, - cvcuda_stream=self.mock_cvcuda_stream, - ) - - with pytest.raises(ValueError, match="Unexpected tensor layout, NCHW expected"): - decoder.generate_decoded_frames() - @patch("nemo_curator.utils.nvcodec_utils.Nvc") @patch("nemo_curator.utils.nvcodec_utils.torch") @patch("nemo_curator.utils.nvcodec_utils.cvcuda")