wip

app/Containerfile

@@ -18,7 +18,17 @@ rm -rf /var/lib/apt/lists/*
 apt-get clean
 
 # Install Python dependencies
-pip3 install onnxruntime_gpu
+case "$(arch)" in
+aarch64)
+  echo "Downloading onnxruntime_gpu from Nvidia"
+  wget -q https://nvidia.box.com/shared/static/i7n40ki3pl2x57vyn4u7e9asyiqlnl7n.whl -O onnxruntime_gpu-1.16.0-cp310-cp310-linux_aarch64.whl
+  pip3 install onnxruntime_gpu-1.16.0-cp310-cp310-linux_aarch64.whl
+  rm -f onnxruntime_gpu-1.16.0-cp310-cp310-linux_aarch64.whl
+  ;;
+x86_64)
+  pip3 install onnxruntime_gpu
+  ;;
+esac
 pip3 install -r requirements.txt
 EOF
 

app/app.py

@@ -82,6 +82,7 @@ if __name__ == "__main__":
     )
     logger = logging.getLogger(__name__)
     ort_sess = ort.InferenceSession(MODEL_PATH, providers=PROVIDERS)
+    logger.info(f"ONNX Runtime device: {ort.get_device()}")
     nparr = np.fromfile(INPUT_IMAGE_PATH, np.uint8)
     nparr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
     preprocessed, scale, original_image = preprocess(nparr)

app/requirements.txt

@@ -1,3 +1,3 @@
 opencv-python-headless
 numpy
-onnxruntime
+onnxruntime_gpu