EdgeAI Jetson Nano LLIE
June 24, 2024 ยท View on GitHub
Introduction
This is a project for Low-light image enhancement with edge AI computations using Jetson Nano.
LLIE approach: CPGA-Net
Python == 3.6
Installation
ONNX
Follow this link Jetson Zoo
- onnxruntime
# Download pip wheel from location above for your version of JetPack
wget https://nvidia.box.com/shared/static/pmsqsiaw4pg9qrbeckcbymho6c01jj4z.whl -O onnxruntime_gpu-1.11.0-cp36-cp36-linux_aarch64.whl
pip3 install onnxruntime_gpu-1.11.0-cp36-cp36-linux_aarch64.whl
TensorRT
TensorRT Installation Guide
TensorRT Intro Notebooks
For installing TensorRT in a conda environment, you can refer to this forum post
- pycuda
- tensorrt
Usage
Pytorch to ONNX
PyTorch tutorial on exporting a simple model to ONNX
Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime
dummy_input = torch.randn(1, 3, 256, 256)
torch.onnx.export(network, dummy_input, MODEL_NAME+'.onnx' )
ONNX
python3 demo_onnx.py
You can try different execution ways, see onnx doc
TRT
trtexec --onnx=OV_enhance_color-llie-ResCBAM_g.onnx --saveEngine=CPGANet_engine.trt --explicitBatch --workspace=128
python3 demo_trt.py
Runtime Performance
Image Resolution: 256x256
| Pytorch | ONNX (CUDA) | ONNX (TRT) | TensorRT |
|---|---|---|---|
| 139.343681 ms | 110.399661 ms | 83.792326 ms | 72.907643 ms |
Efficient Version (DGF)
| Pytorch | ONNX (CUDA) | ONNX (TRT) | TensorRT |
|---|---|---|---|
| 39.909654 ms | 37.896626 ms | 32.811749 ms | ? ms |
Results
Here are some example images enhanced using the LLIE approach:
Pytorch
ONNX
ONNX (TRT)
TensorRT
Image From LIME
Issues
TensorRT doesn't perform normally with unknown issue,