nvImageCodec with TensorFlow¶

[1]:

import os
import cv2
import matplotlib.pyplot as plt
import numpy as np
import cupy as cp
import torch

Set lower log level and import tensorflow

[2]:

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf

Import nvImageCodec

[3]:

from nvidia import nvimgcodec

Setting resource folder

[4]:

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

Created nvImageCodec Decoder

[5]:

decoder = nvimgcodec.Decoder()

Decode JPEG2000 file

[6]:

nv_img = decoder.read(resources_dir + "cat-1046544_640.jp2")

Transfer image to Host memory and display using matplotlib.pyplot

[7]:

plt.imshow(nv_img.cpu())

[7]:

<matplotlib.image.AxesImage at 0x7fae71628b80>

Zero-copy convertion of Image to tensor using DLPack

[8]:

cap = nv_img.to_dlpack()
tf_img = tf.experimental.dlpack.from_dlpack(cap)
print("device:", tf_img.device)
print("dtype of tensor:", tf_img.dtype)

device: /job:localhost/replica:0/task:0/device:GPU:0
dtype of tensor: <dtype: 'uint8'>

Lets flip tensor

[9]:

tf_img_flip = tf.image.flip_up_down(tf_img)

Pass tensor to cuPy and then to numpy to show flipped image

[10]:

cp_img = cp.from_dlpack(tf.experimental.dlpack.to_dlpack(tf_img_flip))
np_img = cp.asnumpy(cp_img)
plt.imshow(np_img)

[10]:

<matplotlib.image.AxesImage at 0x7fae700db490>

Zero-copy convertion of tensor back to Image

[11]:

cap = tf.experimental.dlpack.to_dlpack(tf_img_flip)
nv_flipped_img = nvimgcodec.as_image(cap)

Save as Jpeg2000

[12]:

encoder = nvimgcodec.Encoder()
encoder.write("tf_flipped.j2k", nv_flipped_img)

Load with OpenCV to verify

[13]:

image = cv2.imread("tf_flipped.j2k")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(image)

[13]:

<matplotlib.image.AxesImage at 0x7fae70140370>