From 506ca0d3a49aaa1c0fee4f2620dc9454f4c6ccc1 Mon Sep 17 00:00:00 2001
From: yoni13 <yoni980807@gmail.com>
Date: Sat, 11 Jan 2025 15:47:24 +0800
Subject: [PATCH] Dockerfile for rknn

---
 .github/workflows/docker.yml | 5 ++++-
 machine-learning/Dockerfile  | 6 ++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index ec5c24ba56..1d222d5ddb 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -48,7 +48,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        suffix: ["", "-cuda", "-openvino", "-armnn"]
+        suffix: ["", "-cuda", "-openvino", "-armnn","-rknn"]
     steps:
         - name: Login to GitHub Container Registry
           uses: docker/login-action@v3
@@ -116,6 +116,9 @@ jobs:
           - platforms: linux/arm64
             device: armnn
             suffix: -armnn
+          - platforms: linux/arm64
+            device: rknn
+            suffix: -rknn
 
     steps:
       - name: Checkout
diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile
index 8afe6f0d10..0eb90752f4 100644
--- a/machine-learning/Dockerfile
+++ b/machine-learning/Dockerfile
@@ -15,6 +15,8 @@ RUN mkdir /opt/armnn && \
     cd /opt/ann && \
     sh build.sh
 
+FROM builder-cpu AS builder-rknn
+
 FROM builder-${DEVICE} AS builder
 
 ARG DEVICE
@@ -80,6 +82,10 @@ COPY --from=builder-armnn \
     /opt/ann/build.sh \
     /opt/armnn/
 
+FROM prod-cpu AS prod-rknn
+
+ADD https://github.com/airockchip/rknn-toolkit2/raw/refs/tags/v2.3.0/rknpu2/runtime/Linux/librknn_api/aarch64/librknnrt.so /usr/lib/
+
 FROM prod-${DEVICE} AS prod
 ARG DEVICE