nvImageCodec with cv-cuda (Linux only)

import os
from matplotlib import pyplot as plt
import torch  # import torch before cupy so that we load a newer version of NCCL that torch needs
from torchvision.io.image import read_file, decode_jpeg
from torchvision import models
import cupy as cp
import cvcuda
import numpy as np
import cv2
from nvidia import nvimgcodec

Setting resource folder

resources_dir = os.getenv("PYNVIMGCODEC_EXAMPLES_RESOURCES_DIR", "../assets/images/")

Import nvImageCodec module and create Decoder

decoder = nvimgcodec.Decoder()

Read image with nvImageCodec

inputImage = decoder.read(resources_dir + "tabby_tiger_cat.jpg")
print("size:{}x{}".format(inputImage.width, inputImage.height))

size:720x720

Pass it to cvcuda using as_tensor

nvcvInputTensor = cvcuda.as_tensor(inputImage, "HWC")

Resize with cvcuda

cvcuda_stream = cvcuda.Stream()
with cvcuda_stream:
    nvcvResizeTensor = cvcuda.resize(nvcvInputTensor, (320, 320, 3), cvcuda.Interp.CUBIC)
    nvcvResizeTensor.cuda().__cuda_array_interface__

Write with nvImageCodec

encoder = nvimgcodec.Encoder()
encoder.write("tabby_tiger_cat_320x320.jpg", nvimgcodec.as_image(nvcvResizeTensor.cuda(), cuda_stream = cvcuda_stream.handle))

Verify with OpenCV

image = cv2.imread("tabby_tiger_cat_320x320.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(image)

<matplotlib.image.AxesImage at 0x7f5a37353f40>

Resnet50 classification example from cv-cuda

file_name = resources_dir + "tabby_tiger_cat.jpg"
labelsfile = resources_dir + "../imagenet-classes.txt"

Orignal example code uses torchvision to load image (under the hood it uses nvJpeg)

data = read_file(file_name)
inputImageTmp = decode_jpeg(data, device="cuda")

Now we can change this code to use nvImageCodec and use all formats available with plugins. Please uncomment lines with other images to test it

#file_name = resources_dir + "cat-1046544_640.jp2"
#file_name = resources_dir + "Weimaraner.jpg"

# nvImageCodec has fallback for cpu decoder (only when necessary plugins are installed)
# for codec do not supported yet on GPU so we can read e.g. png
#file_name = resources_dir + "cat-1245673_640.png"

inputImage = decoder.read(file_name)

# A torch tensor/ or nvImageCodec Image can be wrapped into a CVCUDA Object using the "as_tensor"
# function in the specified layout. The datatype and dimensions are derived
# directly from the torch tensor.
nvcvInputTensor = cvcuda.as_tensor(inputImage, "HWC")
image = cp.asnumpy(nvcvInputTensor.cuda())
plt.imshow(image)

#Need 4 dimensions when first is batch size
image_tensors = torch.stack((torch.as_tensor(nvcvInputTensor.cuda()),))
nvcvInputTensor = cvcuda.as_tensor(image_tensors.cuda(), "NHWC")

"""
Preprocessing includes the following sequence of operations.
Resize -> DataType Convert(U8->F32) -> Normalize(Apply mean and std deviation)
-> Interleaved to Planar
"""

# Model settings
layerHeight = 224
layerWidth = 224
batchSize = 1

# Resize
# Resize to the input network dimensions
nvcvResizeTensor = cvcuda.resize(nvcvInputTensor, (1, layerHeight, layerWidth, 3), cvcuda.Interp.CUBIC)

# Convert to the data type and range of values needed by the input layer
# i.e uint8->float. A Scale is applied to normalize the values in the range 0-1
nvcvConvertTensor = cvcuda.convertto(nvcvResizeTensor, np.float32, scale=1 / 255)

"""
The input to the network needs to be normalized based on the mean and
std deviation value to standardize the input data.
"""

# Create a torch tensor to store the mean and standard deviation values for R,G,B
scale = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
scaleTensor = torch.Tensor(scale)
stdTensor = torch.Tensor(std)

# Reshape the the number of channels. The R,G,B values scale and offset will be
# applied to every color plane respectively across the batch
scaleTensor = torch.reshape(scaleTensor, (1, 1, 1, 3)).cuda()
stdTensor = torch.reshape(stdTensor, (1, 1, 1, 3)).cuda()

# Wrap the torch tensor in a CVCUDA Tensor
nvcvScaleTensor = cvcuda.as_tensor(scaleTensor, "NHWC")
nvcvBaseTensor = cvcuda.as_tensor(stdTensor, "NHWC")

# Apply the normalize operator and indicate the scale values are std deviation
# i.e scale = 1/stddev
nvcvNormTensor = cvcuda.normalize(nvcvConvertTensor,
    nvcvBaseTensor, nvcvScaleTensor, cvcuda.NormalizeFlags.SCALE_IS_STDDEV
)

# The final stage in the preprocess pipeline includes converting the RGB buffer
# into a planar buffer
nvcvPreprocessedTensor = cvcuda.reformat(nvcvNormTensor, "NCHW")

# Inference uses pytorch to run a resnet50 model on the preprocessed input and outputs
# the classification scores for 1000 classes
# Load Resnet model pretrained on Imagenet
resnet50 = models.resnet50(pretrained=True)
resnet50.to("cuda")
resnet50.eval()

# Run inference on the preprocessed input
torchPreprocessedTensor = torch.as_tensor(nvcvPreprocessedTensor.cuda(), device="cuda")
inferOutput = resnet50(torchPreprocessedTensor)

"""
Postprocessing function normalizes the classification score from the network and sorts
the scores to get the TopN classification scores.
"""
# top results to print out
topN = 5

# Read and parse the classes
with open(labelsfile, "r") as f:
    classes = [line.strip() for line in f.readlines()]

# Apply softmax to Normalize scores between 0-1
scores = torch.nn.functional.softmax(inferOutput, dim=1)[0]

# Sort output scores in descending order
_, indices = torch.sort(inferOutput, descending=True)

# Display Top N Results
for idx in indices[0][:topN]:
    print("Class : ", classes[idx], " Score : ", scores[idx].item())

/usr/local/lib/python3.9/dist-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.
  warnings.warn(
/usr/local/lib/python3.9/dist-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=ResNet50_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet50_Weights.DEFAULT` to get the most up-to-date weights.
  warnings.warn(msg)
Downloading: "https://download.pytorch.org/models/resnet50-0676ba61.pth" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth
100.0%

Class :  tiger cat  Score :  0.7251120209693909
Class :  tabby, tabby cat  Score :  0.1548738181591034
Class :  Egyptian cat  Score :  0.08538039773702621
Class :  lynx, catamount  Score :  0.02093645930290222
Class :  leopard, Panthera pardus  Score :  0.002835477003827691