updated cudaImage interoperability tests

dusty-nv · Apr 28, 2023 · 09c418e · 09c418e
1 parent 91da77f
commit 09c418e
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 25 deletions.
diff --git a/python/examples/cuda-from-cv.py b/python/examples/cuda-from-cv.py
@@ -40,6 +40,9 @@
 # load the image
 cv_img = cv2.imread(opt.file_in)
 
+if cv_img is None:
+    raise IOError(f"failed to load {file_in}")
+
 print('OpenCV image size: ' + str(cv_img.shape))
 print('OpenCV image type: ' + str(cv_img.dtype))
 

diff --git a/python/examples/cuda-from-numpy.py b/python/examples/cuda-from-numpy.py
@@ -33,7 +33,7 @@
 parser.add_argument("--height", type=int, default=256, help="height of the array (in pixels)")
 parser.add_argument("--depth", type=int, default=4, help="number of color channels in the array (1, 3, or 4)")
 parser.add_argument("--dtype", type=str, default="float32", help="numpy data type: " + " | ".join(sorted({str(key) for key in np.sctypeDict.keys()})))
-parser.add_argument("--filename", type=str, default="cuda-from-numpy.jpg", help="filename of the output test image")
+parser.add_argument("--filename", type=str, default="images/test/cuda-from-numpy.jpg", help="filename of the output test image")
 
 opt = parser.parse_args()
 

diff --git a/python/examples/cuda-from-pytorch.py b/python/examples/cuda-from-pytorch.py
@@ -22,15 +22,21 @@
 #
 
 import argparse
-import torch
+import sys
 
-from jetson_utils import cudaImage
+try:
+    import torch
+except ImportError:
+    print("failed to import torch - if you wish to test PyTorch interoperability, please install it")
+    sys.exit(0)
+
+from jetson_utils import cudaImage, cudaNormalize
 
 
 # parse the command line
 parser = argparse.ArgumentParser('Map cudaImage to PyTorch GPU tensor')
 
-parser.add_argument("--width", type=int, default=8, help="width of the array (in pixels)")
+parser.add_argument("--width", type=int, default=6, help="width of the array (in pixels)")
 parser.add_argument("--height", type=int, default=4, help="height of the array (in pixels)")
 parser.add_argument("--channels", type=int, default=3, help="number of color channels (1, 3, or 4)")
 
@@ -63,16 +69,16 @@ def tensor_image_format(tensor):
 
 
 # allocate a GPU tensor with NCHW layout (strided colors)
-tensor = torch.ones(1, args.channels, args.height, args.width, dtype=torch.float32, device='cuda')
+tensor = torch.rand(1, args.channels, args.height, args.width, dtype=torch.float32, device='cuda')
 
 # transpose the channels to NHWC layout (interleaved colors)
 tensor = tensor.to(memory_format=torch.channels_last)   # or tensor.permute(0, 3, 2, 1)
 
 print("\nPyTorch tensor:")
 print(type(tensor))
-print(hex(tensor.data_ptr()))
-print(tensor.dtype)
-print(tensor.shape)
+print(f"    -- ptr:   {hex(tensor.data_ptr())}")
+print(f"    -- type:  {tensor.dtype}")
+print(f"    -- shape: {tensor.shape}\n")
 print(tensor)
 
 # map to cudaImage using the same underlying memory (any changes will be reflected in the PyTorch tensor)
@@ -81,3 +87,11 @@ def tensor_image_format(tensor):
 print("\ncudaImage:")
 print(cuda_img)
 
+# perform an operation on the cudaImage (scale it by 100x)
+cudaNormalize(cuda_img, (0,1), cuda_img, (0,100))
+
+# print out the PyTorch tensor again to show the values have been updated
+print("\nPyTorch tensor (modified):\n")
+print(tensor)
+
+
diff --git a/python/examples/cuda-to-cv.py b/python/examples/cuda-to-cv.py
@@ -31,7 +31,7 @@
 parser = argparse.ArgumentParser(description='Convert an image from CUDA to OpenCV')
 
 parser.add_argument("file_in", type=str, default="images/jellyfish.jpg", nargs='?', help="filename of the input image to process")
-parser.add_argument("file_out", type=str, default="cuda-to-cv.jpg", nargs='?', help="filename of the output image to save")
+parser.add_argument("file_out", type=str, default="images/test/cuda-to-cv.jpg", nargs='?', help="filename of the output image to save")
 
 opt = parser.parse_args()
 
@@ -53,7 +53,7 @@
 print(bgr_img)
 
 # make sure the GPU is done work before we convert to cv2
-jetson.utils.cudaDeviceSynchronize()
+cudaDeviceSynchronize()
 
 # convert to cv2 image (cv2 images are numpy arrays)
 cv_img = cudaToNumpy(bgr_img)

diff --git a/python/examples/cuda-to-numpy.py b/python/examples/cuda-to-numpy.py
@@ -22,7 +22,7 @@
 #
 
 import argparse
-import numpy
+import numpy as np
 
 from jetson_utils import cudaAllocMapped, cudaToNumpy
 
@@ -32,26 +32,30 @@
 parser.add_argument("--width", type=int, default=4, help="width of the array (in float elements)")
 parser.add_argument("--height", type=int, default=2, help="height of the array (in float elements)")
 
-opt = parser.parse_args()
+args = parser.parse_args()
 
 # allocate cuda memory
-cuda_img = cudaAllocMapped(width=opt.width, height=opt.height, format='rgb32f')
+cuda_img = cudaAllocMapped(width=args.width, height=args.height, format='rgb32f')
 print(cuda_img)
 
-# create a numpy ndarray that references the CUDA memory
-# it won't be copied, but uses the same memory underneath
-array = cudaToNumpy(cuda_img)
+# create a numpy array and do some ops with it
+array = np.ones(cuda_img.shape, np.float32)
 
-print("\ncudaToNumpy() array:")
-print(type(array))
-print(array.dtype)
-print(array.shape)	# numpy dims will be in (height, width, depth) order
-print(array)
+# since cudaImage supports __array__ interface, we can do numpy ops on it directly
+print(np.add(cuda_img, array)) 
+
+# explicitly create a numpy array that references the same CUDA memory (it won't be copied)
+# (this typically isn't necessary to use anymore with cudaImage's __array__ interface)
+mapped_array = cudaToNumpy(cuda_img)
 
-# create another ndarray and do some ops with it
-array2 = numpy.ones(array.shape, numpy.float32)
+print("\ncudaToNumpy() array:")
+print(type(mapped_array))
+print(f"   -- ptr:   {hex(mapped_array.ctypes.data)}")
+print(f"   -- type:  {mapped_array.dtype}")
+print(f"   -- shape: {mapped_array.shape}\n") # numpy dims will be in (height, width, depth) order
+print(mapped_array)                           # this should print out one's
 
 print("\nadding arrays...\n")
-print(array + array2)
+print(array + mapped_array)  # this should print out ones's
 
 
diff --git a/python/examples/cuda-to-pytorch.py b/python/examples/cuda-to-pytorch.py
@@ -22,8 +22,14 @@
 #
 
 import argparse
-import torch
+import sys
 
+try:
+    import torch
+except ImportError:
+    print("failed to import torch - if you wish to test PyTorch interoperability, please install it")
+    sys.exit(0)
+
 from jetson_utils import cudaAllocMapped