Skip to content

Avoid deprecated nvcv_image conversion #1978

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nightcityblade wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Avoid deprecated nvcv_image conversion #1978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a7e5dfc
Avoid deprecated nvcv_image conversion
May 13, 2026
a86ac73
Remove redundant decoded frame layout guard
May 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 12 additions & 17 deletions nemo_curator/utils/nvcodec_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,23 +268,18 @@ 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))
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)
cvcuda_tensor = cvcuda.as_tensor(decoded_frame, "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)

self.local_frame_index = self.local_frame_index + 1
self.decoded_frame_cnt = self.decoded_frame_cnt + 1
Expand Down
58 changes: 4 additions & 54 deletions tests/utils/test_nvcodec_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -714,51 +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_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]

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")
Expand All @@ -780,23 +732,21 @@ 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
mock_cvcuda_tensor = Mock()
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
Expand Down