Merge pull request #209 from Visual-Behavior/load_image

fix image load function for Jetsonnx

Author: thibo73800

alonet/torch2trt/base_exporter.py

@@ -13,7 +13,6 @@
     import onnx_graphsurgeon as gs
     import tensorrt as trt
     import pycuda.driver as cuda
-    from pytorch_quantization import nn as quant_nn
     prod_package_error = None
 except Exception as e:
     prod_package_error = e

aloscene/io/image.py

@@ -1,9 +1,11 @@
-import torchvision
+from aloscene.io.utils.errors import InvalidSampleError
+
+import cv2
 import torch
+import torchvision
+import numpy as np
 from torchvision.io.image import ImageReadMode
 
-from aloscene.io.utils.errors import InvalidSampleError
-
 
 def load_image(image_path):
     """
@@ -22,5 +24,11 @@ def load_image(image_path):
     try:
         image = torchvision.io.read_image(image_path, ImageReadMode.RGB).type(torch.float32)
     except RuntimeError as e:
-        raise InvalidSampleError(f"[Alodataset Warning] Invalid image: {image_path} error={e}")
+        try:
+            image = cv2.imread(image_path)
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            image = np.moveaxis(image, 2, 0)
+            image = torch.Tensor(image).type(torch.float32)
+        except RuntimeError as e:
+            raise InvalidSampleError(f"[Alodataset Warning] Invalid image: {image_path} error={e}")
     return image