nvImageCodec with TensorFlow#
[1]:
import os
import cv2
import matplotlib.pyplot as plt
import cupy as cp
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 0x7fa2b8c3e5c0>
Zero-copy convertion of Image to tensor using DLPack
[8]:
cap = nv_img.to_dlpack()
# tensorflow needs to have all data in the device, as it didn't expose stream which we can synchronize with
cp.cuda.runtime.deviceSynchronize()
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 0x7fa2b4359e70>
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 0x7fa2b43f38b0>
