Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion alonet/torch2trt/base_exporter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
16 changes: 12 additions & 4 deletions aloscene/io/image.py
Original file line number Diff line number Diff line change
@@ -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):
"""
Expand All @@ -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