Initial validated RK3588 face recognition pipeline
This commit is contained in:
commit
e0cce69f15
15 changed files with 10075 additions and 0 deletions
27
.gitignore
vendored
Normal file
27
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*.egg-info/
|
||||
.pytest_cache/
|
||||
|
||||
# Virtual environments
|
||||
.venv/
|
||||
venv/
|
||||
rknn-env/
|
||||
|
||||
# Build artifacts
|
||||
bin/
|
||||
build/
|
||||
*.o
|
||||
*.a
|
||||
*.so
|
||||
|
||||
# Logs / temporary files
|
||||
*.log
|
||||
*.tmp
|
||||
.DS_Store
|
||||
|
||||
# Local/editor files
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
547
README.md
Normal file
547
README.md
Normal file
|
|
@ -0,0 +1,547 @@
|
|||
# Face Recognition on Rockchip RK3588
|
||||
|
||||
Native C++ face detection and face recognition pipeline for Rockchip RK3588 devices using RKNN Runtime.
|
||||
|
||||
Validated on an Orange Pi 5 Plus with Rockchip RK3588.
|
||||
|
||||
## Pipeline
|
||||
|
||||
The pipeline combines:
|
||||
|
||||
- SCRFD 500M KPS face detector at 640x640
|
||||
- 5-point facial landmarks
|
||||
- SFace 2021 face recognition model
|
||||
- RKNN Runtime 2.3.2
|
||||
- Native C++ inference
|
||||
- stb_image for image loading
|
||||
- Cosine similarity for face comparison
|
||||
|
||||
The complete pipeline is:
|
||||
|
||||
```text
|
||||
Image
|
||||
|
|
||||
v
|
||||
stb_image
|
||||
|
|
||||
v
|
||||
SCRFD 500M KPS 640
|
||||
|
|
||||
+--> face bounding boxes
|
||||
|
|
||||
+--> 5 facial landmarks
|
||||
|
|
||||
v
|
||||
SFace alignment
|
||||
112x112
|
||||
|
|
||||
v
|
||||
SFace recognition
|
||||
|
|
||||
v
|
||||
128-D embedding
|
||||
|
|
||||
v
|
||||
L2 normalization
|
||||
|
|
||||
v
|
||||
cosine similarity
|
||||
```
|
||||
|
||||
## Repository Layout
|
||||
|
||||
```text
|
||||
face-rknn-repo/
|
||||
├── src/
|
||||
│ └── face_recognition.cc
|
||||
├── include/
|
||||
│ ├── rknn_api.h
|
||||
│ └── stb_image.h
|
||||
├── models/
|
||||
│ ├── onnx/
|
||||
│ │ ├── SCRFD_500M_KPS_640.onnx
|
||||
│ │ └── face_recognition_sface_2021dec.onnx
|
||||
│ ├── rknn/
|
||||
│ │ ├── SCRFD_500M_KPS_640.rknn
|
||||
│ │ └── face_recognition_sface_2021dec.rknn
|
||||
│ └── SHA256SUMS
|
||||
├── tools/
|
||||
│ └── conversion/
|
||||
│ ├── convert_scrfd_rknn.py
|
||||
│ └── convert_legacy.py
|
||||
├── scripts/
|
||||
│ ├── build.sh
|
||||
│ └── test.sh
|
||||
├── test/
|
||||
│ └── test3f.jpg
|
||||
├── runtime/
|
||||
├── .gitignore
|
||||
└── README.md
|
||||
```
|
||||
|
||||
The runtime directory is intentionally empty in Git. The Rockchip vendor runtime library is installed separately on the target device.
|
||||
|
||||
## Tested Environment
|
||||
|
||||
### Target Device
|
||||
|
||||
- Orange Pi 5 Plus
|
||||
- Rockchip RK3588
|
||||
- ARM64 / aarch64
|
||||
- Linux
|
||||
- RKNN Runtime 2.3.2
|
||||
|
||||
Runtime version:
|
||||
|
||||
```text
|
||||
librknnrt version: 2.3.2
|
||||
```
|
||||
|
||||
Validated runtime library:
|
||||
|
||||
- version: 2.3.2
|
||||
- size: 7,726,232 bytes
|
||||
- MD5: a37ee1d5d664c79836bf6e35b7ef6289
|
||||
|
||||
The runtime library is not committed to this repository.
|
||||
|
||||
## Conversion Environment
|
||||
|
||||
Model conversion was performed on a Debian x86 system using:
|
||||
|
||||
```text
|
||||
Python 3.11.2
|
||||
RKNN Toolkit2 2.3.2
|
||||
RKNN Toolkit2 commit: bd980be9
|
||||
```
|
||||
|
||||
The Python virtual environment used for conversion was:
|
||||
|
||||
```text
|
||||
/home/fabio/photo-ai/rknn-env
|
||||
```
|
||||
|
||||
The virtual environment is not included in the repository.
|
||||
|
||||
# Models
|
||||
|
||||
## SCRFD 500M KPS
|
||||
|
||||
SCRFD is used for face detection and extraction of five facial landmarks.
|
||||
|
||||
ONNX model:
|
||||
|
||||
```text
|
||||
models/onnx/SCRFD_500M_KPS_640.onnx
|
||||
```
|
||||
|
||||
RKNN model:
|
||||
|
||||
```text
|
||||
models/rknn/SCRFD_500M_KPS_640.rknn
|
||||
```
|
||||
|
||||
Target platform:
|
||||
|
||||
```text
|
||||
rk3588
|
||||
```
|
||||
|
||||
Quantization:
|
||||
|
||||
```text
|
||||
disabled
|
||||
```
|
||||
|
||||
Conversion script:
|
||||
|
||||
```text
|
||||
tools/conversion/convert_scrfd_rknn.py
|
||||
```
|
||||
|
||||
## SFace
|
||||
|
||||
SFace is used to generate a 128-dimensional face embedding.
|
||||
|
||||
ONNX model:
|
||||
|
||||
```text
|
||||
models/onnx/face_recognition_sface_2021dec.onnx
|
||||
```
|
||||
|
||||
RKNN model:
|
||||
|
||||
```text
|
||||
models/rknn/face_recognition_sface_2021dec.rknn
|
||||
```
|
||||
|
||||
The RKNN model stored in this repository is the verified model used by the C++ application.
|
||||
|
||||
The exact historical conversion recipe for the SFace RKNN model was not completely preserved, therefore this repository does not claim that the SFace conversion is fully reproducible byte-for-byte.
|
||||
|
||||
# SCRFD Configuration
|
||||
|
||||
The SCRFD input is:
|
||||
|
||||
- 640x640
|
||||
- RGB
|
||||
- FP16
|
||||
- NHWC
|
||||
|
||||
Image preprocessing:
|
||||
|
||||
- top-left letterbox
|
||||
- aspect ratio preserved
|
||||
- padding added to reach 640x640
|
||||
- RGB channel order
|
||||
|
||||
Normalization:
|
||||
|
||||
```text
|
||||
(pixel - 127.5) / 128
|
||||
```
|
||||
|
||||
Detector configuration:
|
||||
|
||||
- strides: 8, 16, 32
|
||||
- anchors per location: 2
|
||||
- detection threshold: 0.50
|
||||
- NMS IoU threshold: 0.45
|
||||
|
||||
The detector produces:
|
||||
|
||||
- bounding boxes
|
||||
- confidence scores
|
||||
- five facial landmarks
|
||||
|
||||
# SFace Configuration
|
||||
|
||||
The five canonical SFace landmarks are:
|
||||
|
||||
```text
|
||||
(38.2946, 51.6963)
|
||||
(73.5318, 51.5014)
|
||||
(56.0252, 71.7366)
|
||||
(41.5493, 92.3655)
|
||||
(70.7299, 92.2041)
|
||||
```
|
||||
|
||||
The detected face is aligned using these landmarks and warped to:
|
||||
|
||||
```text
|
||||
112x112
|
||||
```
|
||||
|
||||
SFace input:
|
||||
|
||||
- RGB
|
||||
- uint8 image values represented as FP16 NHWC
|
||||
- range 0..255
|
||||
- pass_through=1
|
||||
|
||||
The output embedding has 128 dimensions.
|
||||
|
||||
The embedding is L2-normalized before comparison.
|
||||
|
||||
# C++ Application
|
||||
|
||||
The main application is:
|
||||
|
||||
```text
|
||||
src/face_recognition.cc
|
||||
```
|
||||
|
||||
The executable is:
|
||||
|
||||
```text
|
||||
bin/face_recognition
|
||||
```
|
||||
|
||||
The application expects two image paths:
|
||||
|
||||
```text
|
||||
./bin/face_recognition image1.jpg image2.jpg
|
||||
```
|
||||
|
||||
It detects faces in both images, extracts the corresponding embeddings and computes cosine similarity.
|
||||
|
||||
The current test application compares the first detected face in each image.
|
||||
|
||||
# Face Comparison
|
||||
|
||||
Face similarity is computed using cosine similarity between the two L2-normalized 128-dimensional embeddings.
|
||||
|
||||
The comparison threshold used by the current application is:
|
||||
|
||||
```text
|
||||
0.363
|
||||
```
|
||||
|
||||
A similarity above this threshold is considered a match by the current test application.
|
||||
|
||||
This threshold is part of the validated application configuration and should not be interpreted as a universal SFace threshold for every deployment or dataset.
|
||||
|
||||
## Self-Comparison Test
|
||||
|
||||
Comparing an image with itself produces:
|
||||
|
||||
```text
|
||||
cosine similarity = 1.0
|
||||
```
|
||||
|
||||
The reference test image:
|
||||
|
||||
```text
|
||||
test/test3f.jpg
|
||||
```
|
||||
|
||||
contains exactly three detected faces in the validated test.
|
||||
|
||||
# Build
|
||||
|
||||
The application is intended to be compiled on the ARM64/RK3588 target.
|
||||
|
||||
Build script:
|
||||
|
||||
```text
|
||||
scripts/build.sh
|
||||
```
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
./scripts/build.sh
|
||||
```
|
||||
|
||||
The resulting executable is:
|
||||
|
||||
```text
|
||||
bin/face_recognition
|
||||
```
|
||||
|
||||
The executable is linked against the runtime library located in:
|
||||
|
||||
```text
|
||||
runtime/librknnrt.so
|
||||
```
|
||||
|
||||
The build uses an rpath relative to the executable:
|
||||
|
||||
```text
|
||||
$ORIGIN/../runtime
|
||||
```
|
||||
|
||||
This allows the application to use a repository-local runtime without requiring a system-wide installation.
|
||||
|
||||
# Test
|
||||
|
||||
The test script is:
|
||||
|
||||
```text
|
||||
scripts/test.sh
|
||||
```
|
||||
|
||||
Run the default self-comparison:
|
||||
|
||||
```bash
|
||||
./scripts/test.sh
|
||||
```
|
||||
|
||||
This uses:
|
||||
|
||||
```text
|
||||
test/test3f.jpg
|
||||
```
|
||||
|
||||
for both inputs.
|
||||
|
||||
Two explicit images can also be supplied:
|
||||
|
||||
```bash
|
||||
./scripts/test.sh image1.jpg image2.jpg
|
||||
```
|
||||
|
||||
# Model Integrity
|
||||
|
||||
SHA256 checksums for all committed models are stored in:
|
||||
|
||||
```text
|
||||
models/SHA256SUMS
|
||||
```
|
||||
|
||||
Verify the models with:
|
||||
|
||||
```bash
|
||||
cd models
|
||||
sha256sum -c SHA256SUMS
|
||||
```
|
||||
|
||||
Expected result:
|
||||
|
||||
```text
|
||||
face_recognition_sface_2021dec.onnx: OK
|
||||
SCRFD_500M_KPS_640.onnx: OK
|
||||
face_recognition_sface_2021dec.rknn: OK
|
||||
SCRFD_500M_KPS_640.rknn: OK
|
||||
```
|
||||
|
||||
Current SHA256 values:
|
||||
|
||||
```text
|
||||
face_recognition_sface_2021dec.onnx
|
||||
0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79
|
||||
|
||||
SCRFD_500M_KPS_640.onnx
|
||||
857efab2e0a5184ec86ffa7d0bf33ac94da92591e7650d5353622ce367218faf
|
||||
|
||||
face_recognition_sface_2021dec.rknn
|
||||
5f36840c6fea8a4772a45fe5eb3456b7bd2031e2986947154fcbd165f918f2ec
|
||||
|
||||
SCRFD_500M_KPS_640.rknn
|
||||
7d74abdedebc5fe25c98195db75cc915df74dd58d941be60c1c260186ac764e2
|
||||
```
|
||||
|
||||
# Conversion
|
||||
|
||||
Model conversion was performed separately from the target runtime.
|
||||
|
||||
The repository contains the conversion scripts used for the validated SCRFD conversion and the historical generic conversion tooling.
|
||||
|
||||
## SCRFD Conversion
|
||||
|
||||
The SCRFD conversion script is:
|
||||
|
||||
```text
|
||||
tools/conversion/convert_scrfd_rknn.py
|
||||
```
|
||||
|
||||
Its essential configuration is:
|
||||
|
||||
```python
|
||||
from rknn.api import RKNN
|
||||
|
||||
ONNX_MODEL = "SCRFD_500M_KPS_640.onnx"
|
||||
RKNN_MODEL = "SCRFD_500M_KPS_640.rknn"
|
||||
|
||||
rknn = RKNN(verbose=True)
|
||||
rknn.config(target_platform="rk3588")
|
||||
rknn.load_onnx(model=ONNX_MODEL)
|
||||
rknn.build(do_quantization=False)
|
||||
rknn.export_rknn(RKNN_MODEL)
|
||||
rknn.release()
|
||||
```
|
||||
|
||||
The script expects the ONNX model in the current working directory.
|
||||
|
||||
## SFace Conversion
|
||||
|
||||
The repository includes:
|
||||
|
||||
```text
|
||||
tools/conversion/convert_legacy.py
|
||||
```
|
||||
|
||||
This is historical generic ONNX-to-RKNN conversion tooling.
|
||||
|
||||
The validated SFace RKNN model is committed to the repository, but the complete original conversion procedure, including all intermediate optimization steps and exact conversion inputs, was not fully preserved.
|
||||
|
||||
Therefore:
|
||||
|
||||
- the committed SFace RKNN model is reproducible as an artifact
|
||||
- its SHA256 checksum is verified
|
||||
- the exact original byte-for-byte conversion process is not claimed to be reproducible
|
||||
|
||||
# Reproducibility
|
||||
|
||||
The repository is intended to preserve the working state of the validated pipeline.
|
||||
|
||||
The following are versioned:
|
||||
|
||||
- C++ source
|
||||
- RKNN API header
|
||||
- stb_image header
|
||||
- ONNX models
|
||||
- RKNN models
|
||||
- conversion scripts
|
||||
- build script
|
||||
- test script
|
||||
- test image
|
||||
- SHA256 checksums
|
||||
- documentation
|
||||
|
||||
The following are intentionally not versioned:
|
||||
|
||||
- Python virtual environments
|
||||
- build artifacts
|
||||
- compiled executables
|
||||
- shared libraries
|
||||
- vendor runtime binaries
|
||||
- temporary conversion files
|
||||
- editor configuration
|
||||
|
||||
The target application can therefore be rebuilt on an ARM64/RK3588 system while keeping the validated model artifacts and source code under version control.
|
||||
|
||||
# Runtime Library
|
||||
|
||||
The Rockchip RKNN runtime is a vendor-provided binary.
|
||||
|
||||
The validated version is:
|
||||
|
||||
```text
|
||||
2.3.2
|
||||
```
|
||||
|
||||
The repository deliberately does not commit:
|
||||
|
||||
```text
|
||||
librknnrt.so
|
||||
```
|
||||
|
||||
The target device must provide a compatible RKNN Runtime installation or the runtime library must be placed locally in:
|
||||
|
||||
```text
|
||||
runtime/librknnrt.so
|
||||
```
|
||||
|
||||
The build system uses that local library when compiling.
|
||||
|
||||
# Third-Party Components
|
||||
|
||||
This project uses third-party components including:
|
||||
|
||||
- Rockchip RKNN Runtime
|
||||
- Rockchip RKNN Toolkit2
|
||||
- SCRFD
|
||||
- SFace
|
||||
- stb_image
|
||||
|
||||
Their respective licenses and redistribution terms remain applicable.
|
||||
|
||||
This repository does not claim ownership of those third-party components.
|
||||
|
||||
# License
|
||||
|
||||
The application source in this repository should be considered project-specific code.
|
||||
|
||||
Third-party components, models, headers and runtime libraries remain subject to their original licenses and terms.
|
||||
|
||||
Before redistributing the complete repository or its models, verify the applicable licenses and redistribution permissions for each third-party component.
|
||||
|
||||
# Status
|
||||
|
||||
Current validated status:
|
||||
|
||||
- SCRFD RKNN inference: working
|
||||
- SFace RKNN inference: working
|
||||
- Face landmark extraction: working
|
||||
- Face alignment: working
|
||||
- 128-D embedding generation: working
|
||||
- L2 normalization: working
|
||||
- Cosine similarity: working
|
||||
- Self-comparison: cosine similarity 1.0
|
||||
- Three-face test image: validated
|
||||
- ARM64/RK3588 native C++ application: working
|
||||
- Local RKNN runtime loading: working
|
||||
- Model SHA256 verification: working
|
||||
|
||||
The repository represents the validated working baseline of the RK3588 face-recognition pipeline.
|
||||
804
include/rknn_api.h
Normal file
804
include/rknn_api.h
Normal file
|
|
@ -0,0 +1,804 @@
|
|||
/****************************************************************************
|
||||
*
|
||||
* Copyright (c) 2017 - 2022 by Rockchip Corp. All rights reserved.
|
||||
*
|
||||
* The material in this file is confidential and contains trade secrets
|
||||
* of Rockchip Corporation. This is proprietary information owned by
|
||||
* Rockchip Corporation. No part of this work may be disclosed,
|
||||
* reproduced, copied, transmitted, or used in any way for any purpose,
|
||||
* without the express written permission of Rockchip Corporation.
|
||||
*
|
||||
*****************************************************************************/
|
||||
|
||||
|
||||
#ifndef _RKNN_API_H
|
||||
#define _RKNN_API_H
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
/*
|
||||
Definition of extended flag for rknn_init.
|
||||
*/
|
||||
/* set high priority context. */
|
||||
#define RKNN_FLAG_PRIOR_HIGH 0x00000000
|
||||
|
||||
/* set medium priority context */
|
||||
#define RKNN_FLAG_PRIOR_MEDIUM 0x00000001
|
||||
|
||||
/* set low priority context. */
|
||||
#define RKNN_FLAG_PRIOR_LOW 0x00000002
|
||||
|
||||
/* asynchronous mode.
|
||||
when enable, rknn_outputs_get will not block for too long because it directly retrieves the result of
|
||||
the previous frame which can increase the frame rate on single-threaded mode, but at the cost of
|
||||
rknn_outputs_get not retrieves the result of the current frame.
|
||||
in multi-threaded mode you do not need to turn this mode on. */
|
||||
#define RKNN_FLAG_ASYNC_MASK 0x00000004
|
||||
|
||||
/* collect performance mode.
|
||||
when enable, you can get detailed performance reports via rknn_query(ctx, RKNN_QUERY_PERF_DETAIL, ...),
|
||||
but it will reduce the frame rate. */
|
||||
#define RKNN_FLAG_COLLECT_PERF_MASK 0x00000008
|
||||
|
||||
/* allocate all memory in outside, includes weight/internal/inputs/outputs */
|
||||
#define RKNN_FLAG_MEM_ALLOC_OUTSIDE 0x00000010
|
||||
|
||||
/* weight sharing with the same network structure */
|
||||
#define RKNN_FLAG_SHARE_WEIGHT_MEM 0x00000020
|
||||
|
||||
/* send fence fd from outside */
|
||||
#define RKNN_FLAG_FENCE_IN_OUTSIDE 0x00000040
|
||||
|
||||
/* get fence fd from inside */
|
||||
#define RKNN_FLAG_FENCE_OUT_OUTSIDE 0x00000080
|
||||
|
||||
/* dummy init flag: could only get total_weight_size and total_internal_size by rknn_query*/
|
||||
#define RKNN_FLAG_COLLECT_MODEL_INFO_ONLY 0x00000100
|
||||
|
||||
/* allocate internal memory in outside */
|
||||
#define RKNN_FLAG_INTERNAL_ALLOC_OUTSIDE 0x00000200
|
||||
|
||||
/* set GPU as the preferred execution backend When the operator is not supported by the NPU */
|
||||
#define RKNN_FLAG_EXECUTE_FALLBACK_PRIOR_DEVICE_GPU 0x00000400
|
||||
|
||||
/* enable allocate sram type buffers */
|
||||
#define RKNN_FLAG_ENABLE_SRAM 0x00000800
|
||||
|
||||
/* sram type buffers are shared among different contexts */
|
||||
#define RKNN_FLAG_SHARE_SRAM 0x00001000
|
||||
|
||||
/* default nice -19, this flag can disable default priority */
|
||||
#define RKNN_FLAG_DISABLE_PROC_HIGH_PRIORITY 0x00002000
|
||||
|
||||
/* don't flush input buffer cache, the user must ensure that the input tensor has flushed the cache before calling rknn_run.
|
||||
!!! Don't use this flags when you call rknn_inputs_set() to set input data. */
|
||||
#define RKNN_FLAG_DISABLE_FLUSH_INPUT_MEM_CACHE 0x00004000
|
||||
|
||||
/* Don't invalid output buffer cache.
|
||||
Users cannot directly access output_mem->virt_addr,
|
||||
which will cause cache consistency problems.
|
||||
If you want to use output_mem->virt_addr,
|
||||
you must use rknn_mem_sync (ctx, mem, RKNN_MEMORY_SYNC_FROM_DEVICE) to flush the cache.
|
||||
This flags is generally used when the output data of the NPU is not accessed by the CPU,
|
||||
but is accessed by the GPU or RGA to reduce the time required to flush the cache.
|
||||
!!! Don't use this flags when you call rknn_outputs_get() to get output data.*/
|
||||
#define RKNN_FLAG_DISABLE_FLUSH_OUTPUT_MEM_CACHE 0x00008000
|
||||
|
||||
/* This flag is used when the model data buffer is allocated by NPU, and can be accessed by NPU directly. */
|
||||
#define RKNN_FLAG_MODEL_BUFFER_ZERO_COPY 0x00010000
|
||||
|
||||
/* This flag is a memory allocation flag, which is used in rknn_create_mem2() when no context is available. */
|
||||
#define RKNN_MEM_FLAG_ALLOC_NO_CONTEXT 0x00020000
|
||||
|
||||
|
||||
/*
|
||||
Error code returned by the RKNN API.
|
||||
*/
|
||||
#define RKNN_SUCC 0 /* execute succeed. */
|
||||
#define RKNN_ERR_FAIL -1 /* execute failed. */
|
||||
#define RKNN_ERR_TIMEOUT -2 /* execute timeout. */
|
||||
#define RKNN_ERR_DEVICE_UNAVAILABLE -3 /* device is unavailable. */
|
||||
#define RKNN_ERR_MALLOC_FAIL -4 /* memory malloc fail. */
|
||||
#define RKNN_ERR_PARAM_INVALID -5 /* parameter is invalid. */
|
||||
#define RKNN_ERR_MODEL_INVALID -6 /* model is invalid. */
|
||||
#define RKNN_ERR_CTX_INVALID -7 /* context is invalid. */
|
||||
#define RKNN_ERR_INPUT_INVALID -8 /* input is invalid. */
|
||||
#define RKNN_ERR_OUTPUT_INVALID -9 /* output is invalid. */
|
||||
#define RKNN_ERR_DEVICE_UNMATCH -10 /* the device is unmatch, please update rknn sdk
|
||||
and npu driver/firmware. */
|
||||
#define RKNN_ERR_INCOMPATILE_PRE_COMPILE_MODEL -11 /* This RKNN model use pre_compile mode, but not compatible with current driver. */
|
||||
#define RKNN_ERR_INCOMPATILE_OPTIMIZATION_LEVEL_VERSION -12 /* This RKNN model set optimization level, but not compatible with current driver. */
|
||||
#define RKNN_ERR_TARGET_PLATFORM_UNMATCH -13 /* This RKNN model set target platform, but not compatible with current platform. */
|
||||
|
||||
/*
|
||||
Definition for tensor
|
||||
*/
|
||||
#define RKNN_MAX_DIMS 16 /* maximum dimension of tensor. */
|
||||
#define RKNN_MAX_NUM_CHANNEL 15 /* maximum channel number of input tensor. */
|
||||
#define RKNN_MAX_NAME_LEN 256 /* maximum name lenth of tensor. */
|
||||
#define RKNN_MAX_DYNAMIC_SHAPE_NUM 512 /* maximum number of dynamic shape for each input. */
|
||||
|
||||
#ifdef __arm__
|
||||
typedef uint32_t rknn_context;
|
||||
#else
|
||||
typedef uint64_t rknn_context;
|
||||
#endif
|
||||
|
||||
|
||||
/*
|
||||
The query command for rknn_query
|
||||
*/
|
||||
typedef enum _rknn_query_cmd {
|
||||
RKNN_QUERY_IN_OUT_NUM = 0, /* query the number of input & output tensor. */
|
||||
RKNN_QUERY_INPUT_ATTR = 1, /* query the attribute of input tensor. */
|
||||
RKNN_QUERY_OUTPUT_ATTR = 2, /* query the attribute of output tensor. */
|
||||
RKNN_QUERY_PERF_DETAIL = 3, /* query the detail performance, need set
|
||||
RKNN_FLAG_COLLECT_PERF_MASK when call rknn_init,
|
||||
this query needs to be valid after rknn_outputs_get. */
|
||||
RKNN_QUERY_PERF_RUN = 4, /* query the time of run,
|
||||
this query needs to be valid after rknn_outputs_get. */
|
||||
RKNN_QUERY_SDK_VERSION = 5, /* query the sdk & driver version */
|
||||
|
||||
RKNN_QUERY_MEM_SIZE = 6, /* query the weight & internal memory size */
|
||||
RKNN_QUERY_CUSTOM_STRING = 7, /* query the custom string */
|
||||
|
||||
RKNN_QUERY_NATIVE_INPUT_ATTR = 8, /* query the attribute of native input tensor. */
|
||||
RKNN_QUERY_NATIVE_OUTPUT_ATTR = 9, /* query the attribute of native output tensor. */
|
||||
|
||||
RKNN_QUERY_NATIVE_NC1HWC2_INPUT_ATTR = 8, /* query the attribute of native input tensor. */
|
||||
RKNN_QUERY_NATIVE_NC1HWC2_OUTPUT_ATTR = 9, /* query the attribute of native output tensor. */
|
||||
|
||||
RKNN_QUERY_NATIVE_NHWC_INPUT_ATTR = 10, /* query the attribute of native input tensor. */
|
||||
RKNN_QUERY_NATIVE_NHWC_OUTPUT_ATTR = 11, /* query the attribute of native output tensor. */
|
||||
|
||||
RKNN_QUERY_DEVICE_MEM_INFO = 12, /* query the attribute of rknn memory information. */
|
||||
|
||||
RKNN_QUERY_INPUT_DYNAMIC_RANGE = 13, /* query the dynamic shape range of rknn input tensor. */
|
||||
RKNN_QUERY_CURRENT_INPUT_ATTR = 14, /* query the current shape of rknn input tensor, only valid for dynamic rknn model*/
|
||||
RKNN_QUERY_CURRENT_OUTPUT_ATTR = 15, /* query the current shape of rknn output tensor, only valid for dynamic rknn model*/
|
||||
|
||||
RKNN_QUERY_CURRENT_NATIVE_INPUT_ATTR = 16, /* query the current native shape of rknn input tensor, only valid for dynamic rknn model*/
|
||||
RKNN_QUERY_CURRENT_NATIVE_OUTPUT_ATTR = 17, /* query the current native shape of rknn output tensor, only valid for dynamic rknn model*/
|
||||
|
||||
|
||||
RKNN_QUERY_CMD_MAX
|
||||
} rknn_query_cmd;
|
||||
|
||||
/*
|
||||
the tensor data type.
|
||||
*/
|
||||
typedef enum _rknn_tensor_type {
|
||||
RKNN_TENSOR_FLOAT32 = 0, /* data type is float32. */
|
||||
RKNN_TENSOR_FLOAT16, /* data type is float16. */
|
||||
RKNN_TENSOR_INT8, /* data type is int8. */
|
||||
RKNN_TENSOR_UINT8, /* data type is uint8. */
|
||||
RKNN_TENSOR_INT16, /* data type is int16. */
|
||||
RKNN_TENSOR_UINT16, /* data type is uint16. */
|
||||
RKNN_TENSOR_INT32, /* data type is int32. */
|
||||
RKNN_TENSOR_UINT32, /* data type is uint32. */
|
||||
RKNN_TENSOR_INT64, /* data type is int64. */
|
||||
RKNN_TENSOR_BOOL,
|
||||
RKNN_TENSOR_INT4,
|
||||
RKNN_TENSOR_BFLOAT16,
|
||||
|
||||
RKNN_TENSOR_TYPE_MAX
|
||||
} rknn_tensor_type;
|
||||
|
||||
inline static const char* get_type_string(rknn_tensor_type type)
|
||||
{
|
||||
switch(type) {
|
||||
case RKNN_TENSOR_FLOAT32: return "FP32";
|
||||
case RKNN_TENSOR_FLOAT16: return "FP16";
|
||||
case RKNN_TENSOR_INT8: return "INT8";
|
||||
case RKNN_TENSOR_UINT8: return "UINT8";
|
||||
case RKNN_TENSOR_INT16: return "INT16";
|
||||
case RKNN_TENSOR_UINT16: return "UINT16";
|
||||
case RKNN_TENSOR_INT32: return "INT32";
|
||||
case RKNN_TENSOR_UINT32: return "UINT32";
|
||||
case RKNN_TENSOR_INT64: return "INT64";
|
||||
case RKNN_TENSOR_BOOL: return "BOOL";
|
||||
case RKNN_TENSOR_INT4: return "INT4";
|
||||
case RKNN_TENSOR_BFLOAT16: return "BF16";
|
||||
default: return "UNKNOW";
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
the quantitative type.
|
||||
*/
|
||||
typedef enum _rknn_tensor_qnt_type {
|
||||
RKNN_TENSOR_QNT_NONE = 0, /* none. */
|
||||
RKNN_TENSOR_QNT_DFP, /* dynamic fixed point. */
|
||||
RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC, /* asymmetric affine. */
|
||||
|
||||
RKNN_TENSOR_QNT_MAX
|
||||
} rknn_tensor_qnt_type;
|
||||
|
||||
inline static const char* get_qnt_type_string(rknn_tensor_qnt_type type)
|
||||
{
|
||||
switch(type) {
|
||||
case RKNN_TENSOR_QNT_NONE: return "NONE";
|
||||
case RKNN_TENSOR_QNT_DFP: return "DFP";
|
||||
case RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC: return "AFFINE";
|
||||
default: return "UNKNOW";
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
the tensor data format.
|
||||
*/
|
||||
typedef enum _rknn_tensor_format {
|
||||
RKNN_TENSOR_NCHW = 0, /* data format is NCHW. */
|
||||
RKNN_TENSOR_NHWC, /* data format is NHWC. */
|
||||
RKNN_TENSOR_NC1HWC2, /* data format is NC1HWC2. */
|
||||
RKNN_TENSOR_UNDEFINED,
|
||||
|
||||
RKNN_TENSOR_FORMAT_MAX
|
||||
} rknn_tensor_format;
|
||||
|
||||
/*
|
||||
the mode of running on target NPU core.
|
||||
*/
|
||||
typedef enum _rknn_core_mask {
|
||||
RKNN_NPU_CORE_AUTO = 0, /* default, run on NPU core randomly. */
|
||||
RKNN_NPU_CORE_0 = 1, /* run on NPU core 0. */
|
||||
RKNN_NPU_CORE_1 = 2, /* run on NPU core 1. */
|
||||
RKNN_NPU_CORE_2 = 4, /* run on NPU core 2. */
|
||||
RKNN_NPU_CORE_0_1 = RKNN_NPU_CORE_0 | RKNN_NPU_CORE_1, /* run on NPU core 0 and core 1. */
|
||||
RKNN_NPU_CORE_0_1_2 = RKNN_NPU_CORE_0_1 | RKNN_NPU_CORE_2, /* run on NPU core 0 and core 1 and core 2. */
|
||||
RKNN_NPU_CORE_ALL = 0xffff, /* auto choice, run on NPU cores depending on platform */
|
||||
|
||||
RKNN_NPU_CORE_UNDEFINED,
|
||||
} rknn_core_mask;
|
||||
|
||||
inline static const char* get_format_string(rknn_tensor_format fmt)
|
||||
{
|
||||
switch(fmt) {
|
||||
case RKNN_TENSOR_NCHW: return "NCHW";
|
||||
case RKNN_TENSOR_NHWC: return "NHWC";
|
||||
case RKNN_TENSOR_NC1HWC2: return "NC1HWC2";
|
||||
case RKNN_TENSOR_UNDEFINED: return "UNDEFINED";
|
||||
default: return "UNKNOW";
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
the information for RKNN_QUERY_IN_OUT_NUM.
|
||||
*/
|
||||
typedef struct _rknn_input_output_num {
|
||||
uint32_t n_input; /* the number of input. */
|
||||
uint32_t n_output; /* the number of output. */
|
||||
} rknn_input_output_num;
|
||||
|
||||
/*
|
||||
the information for RKNN_QUERY_INPUT_ATTR / RKNN_QUERY_OUTPUT_ATTR.
|
||||
*/
|
||||
typedef struct _rknn_tensor_attr {
|
||||
uint32_t index; /* input parameter, the index of input/output tensor,
|
||||
need set before call rknn_query. */
|
||||
|
||||
uint32_t n_dims; /* the number of dimensions. */
|
||||
uint32_t dims[RKNN_MAX_DIMS]; /* the dimensions array. */
|
||||
char name[RKNN_MAX_NAME_LEN]; /* the name of tensor. */
|
||||
|
||||
uint32_t n_elems; /* the number of elements. */
|
||||
uint32_t size; /* the bytes size of tensor. */
|
||||
|
||||
rknn_tensor_format fmt; /* the data format of tensor. */
|
||||
rknn_tensor_type type; /* the data type of tensor. */
|
||||
rknn_tensor_qnt_type qnt_type; /* the quantitative type of tensor. */
|
||||
int8_t fl; /* fractional length for RKNN_TENSOR_QNT_DFP. */
|
||||
int32_t zp; /* zero point for RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC. */
|
||||
float scale; /* scale for RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC. */
|
||||
|
||||
uint32_t w_stride; /* the stride of tensor along the width dimention of input,
|
||||
Note: it is read-only, 0 means equal to width. */
|
||||
uint32_t size_with_stride; /* the bytes size of tensor with stride. */
|
||||
|
||||
uint8_t pass_through; /* pass through mode, for rknn_set_io_mem interface.
|
||||
if TRUE, the buf data is passed directly to the input node of the rknn model
|
||||
without any conversion. the following variables do not need to be set.
|
||||
if FALSE, the buf data is converted into an input consistent with the model
|
||||
according to the following type and fmt. so the following variables
|
||||
need to be set.*/
|
||||
uint32_t h_stride; /* the stride along the height dimention of input,
|
||||
Note: it is write-only, if it was set to 0, h_stride = height. */
|
||||
} rknn_tensor_attr;
|
||||
|
||||
typedef struct _rknn_input_range {
|
||||
uint32_t index; /* input parameter, the index of input/output tensor,
|
||||
need set before call rknn_query. */
|
||||
uint32_t shape_number; /* the number of shape. */
|
||||
rknn_tensor_format fmt; /* the data format of tensor. */
|
||||
char name[RKNN_MAX_NAME_LEN]; /* the name of tensor. */
|
||||
uint32_t dyn_range[RKNN_MAX_DYNAMIC_SHAPE_NUM][RKNN_MAX_DIMS]; /* the dynamic input dimensions range. */
|
||||
uint32_t n_dims; /* the number of dimensions. */
|
||||
|
||||
} rknn_input_range;
|
||||
|
||||
/*
|
||||
the information for RKNN_QUERY_PERF_DETAIL.
|
||||
*/
|
||||
typedef struct _rknn_perf_detail {
|
||||
char* perf_data; /* the string pointer of perf detail. don't need free it by user. */
|
||||
uint64_t data_len; /* the string length. */
|
||||
} rknn_perf_detail;
|
||||
|
||||
/*
|
||||
the information for RKNN_QUERY_PERF_RUN.
|
||||
*/
|
||||
typedef struct _rknn_perf_run {
|
||||
int64_t run_duration; /* real inference time (us) */
|
||||
} rknn_perf_run;
|
||||
|
||||
/*
|
||||
the information for RKNN_QUERY_SDK_VERSION.
|
||||
*/
|
||||
typedef struct _rknn_sdk_version {
|
||||
char api_version[256]; /* the version of rknn api. */
|
||||
char drv_version[256]; /* the version of rknn driver. */
|
||||
} rknn_sdk_version;
|
||||
|
||||
/*
|
||||
the information for RKNN_QUERY_MEM_SIZE.
|
||||
*/
|
||||
typedef struct _rknn_mem_size {
|
||||
uint32_t total_weight_size; /* the weight memory size */
|
||||
uint32_t total_internal_size; /* the internal memory size, exclude inputs/outputs */
|
||||
uint64_t total_dma_allocated_size; /* total dma memory allocated size */
|
||||
uint32_t total_sram_size; /* total system sram size reserved for rknn */
|
||||
uint32_t free_sram_size; /* free system sram size reserved for rknn */
|
||||
uint32_t reserved[10]; /* reserved */
|
||||
} rknn_mem_size;
|
||||
|
||||
/*
|
||||
the information for RKNN_QUERY_CUSTOM_STRING.
|
||||
*/
|
||||
typedef struct _rknn_custom_string {
|
||||
char string[1024]; /* the string of custom, lengths max to 1024 bytes */
|
||||
} rknn_custom_string;
|
||||
|
||||
/*
|
||||
The flags of rknn_tensor_mem.
|
||||
*/
|
||||
typedef enum _rknn_tensor_mem_flags {
|
||||
RKNN_TENSOR_MEMORY_FLAGS_ALLOC_INSIDE = 1, /*Used to mark in rknn_destroy_mem() whether it is necessary to release the "mem" pointer itself.
|
||||
If the flag RKNN_TENSOR_MEMORY_FLAGS_ALLOC_INSIDE is set, rknn_destroy_mem() will call free(mem).*/
|
||||
RKNN_TENSOR_MEMORY_FLAGS_FROM_FD = 2, /*Used to mark in rknn_create_mem_from_fd() whether it is necessary to release the "mem" pointer itself.
|
||||
If the flag RKNN_TENSOR_MEMORY_FLAGS_FROM_FD is set, rknn_destroy_mem() will call free(mem).*/
|
||||
RKNN_TENSOR_MEMORY_FLAGS_FROM_PHYS = 3, /*Used to mark in rknn_create_mem_from_phys() whether it is necessary to release the "mem" pointer itself.
|
||||
If the flag RKNN_TENSOR_MEMORY_FLAGS_FROM_PHYS is set, rknn_destroy_mem() will call free(mem).*/
|
||||
RKNN_TENSOR_MEMORY_FLAGS_UNKNOWN
|
||||
} rknn_tensor_mem_flags;
|
||||
|
||||
/*
|
||||
The mode to allocate rknn memory.
|
||||
*/
|
||||
typedef enum _rknn_mem_alloc_flags {
|
||||
RKNN_FLAG_MEMORY_FLAGS_DEFAULT = 0 << 0, /* Same with RKNN_FLAG_MEMORY_CACHEABLE */
|
||||
RKNN_FLAG_MEMORY_CACHEABLE = 1 << 0, /* Create Cacheable memory. */
|
||||
RKNN_FLAG_MEMORY_NON_CACHEABLE = 1 << 1, /* Create NON-Cacheable memory. */
|
||||
RKNN_FLAG_MEMORY_TRY_ALLOC_SRAM = 1 << 2, /* Try to allocate memory in SRAM if possible. if SRAM is not enough, allocate rest memory in DRAM. */
|
||||
} rknn_mem_alloc_flags;
|
||||
|
||||
/*
|
||||
The mode to sync cacheable rknn memory.
|
||||
*/
|
||||
typedef enum _rknn_mem_sync_mode {
|
||||
RKNN_MEMORY_SYNC_TO_DEVICE = 0x1, /* the mode used for consistency of device access after CPU accesses data. */
|
||||
RKNN_MEMORY_SYNC_FROM_DEVICE = 0x2, /* the mode used for consistency of CPU access after device accesses data. */
|
||||
RKNN_MEMORY_SYNC_BIDIRECTIONAL = RKNN_MEMORY_SYNC_TO_DEVICE | RKNN_MEMORY_SYNC_FROM_DEVICE, /* the mode used for consistency of data access
|
||||
between device and CPU in both directions. */
|
||||
} rknn_mem_sync_mode;
|
||||
|
||||
/*
|
||||
the memory information of tensor.
|
||||
*/
|
||||
typedef struct _rknn_tensor_memory {
|
||||
void* virt_addr; /* the virtual address of tensor buffer. */
|
||||
uint64_t phys_addr; /* the physical address of tensor buffer. */
|
||||
int32_t fd; /* the fd of tensor buffer. */
|
||||
int32_t offset; /* indicates the offset of the memory. */
|
||||
uint32_t size; /* the size of tensor buffer. */
|
||||
uint32_t flags; /* the flags of tensor buffer, reserved */
|
||||
void * priv_data; /* the private data of tensor buffer. */
|
||||
} rknn_tensor_mem;
|
||||
|
||||
/*
|
||||
the input information for rknn_input_set.
|
||||
*/
|
||||
typedef struct _rknn_input {
|
||||
uint32_t index; /* the input index. */
|
||||
void* buf; /* the input buf for index. */
|
||||
uint32_t size; /* the size of input buf. */
|
||||
uint8_t pass_through; /* pass through mode.
|
||||
if TRUE, the buf data is passed directly to the input node of the rknn model
|
||||
without any conversion. the following variables do not need to be set.
|
||||
if FALSE, the buf data is converted into an input consistent with the model
|
||||
according to the following type and fmt. so the following variables
|
||||
need to be set.*/
|
||||
rknn_tensor_type type; /* the data type of input buf. */
|
||||
rknn_tensor_format fmt; /* the data format of input buf.
|
||||
currently the internal input format of NPU is NCHW by default.
|
||||
so entering NCHW data can avoid the format conversion in the driver. */
|
||||
} rknn_input;
|
||||
|
||||
/*
|
||||
the output information for rknn_outputs_get.
|
||||
*/
|
||||
typedef struct _rknn_output {
|
||||
uint8_t want_float; /* want transfer output data to float */
|
||||
uint8_t is_prealloc; /* whether buf is pre-allocated.
|
||||
if TRUE, the following variables need to be set.
|
||||
if FALSE, the following variables do not need to be set. */
|
||||
uint32_t index; /* the output index. */
|
||||
void* buf; /* the output buf for index.
|
||||
when is_prealloc = FALSE and rknn_outputs_release called,
|
||||
this buf pointer will be free and don't use it anymore. */
|
||||
uint32_t size; /* the size of output buf. */
|
||||
} rknn_output;
|
||||
|
||||
/*
|
||||
the extend information for rknn_init.
|
||||
*/
|
||||
typedef struct _rknn_init_extend {
|
||||
rknn_context ctx; /* rknn context */
|
||||
int32_t real_model_offset; /* real rknn model file offset, only valid when init context with rknn file path and zero-copy model model */
|
||||
uint32_t real_model_size; /* real rknn model file size, only valid when init context with rknn file path and zero-copy model model */
|
||||
int32_t model_buffer_fd; /* the fd of model buffer. */
|
||||
uint32_t model_buffer_flags; /* the flags of model_buffer */
|
||||
uint8_t reserved[112]; /* reserved */
|
||||
} rknn_init_extend;
|
||||
|
||||
/*
|
||||
the extend information for rknn_run.
|
||||
*/
|
||||
typedef struct _rknn_run_extend {
|
||||
uint64_t frame_id; /* output parameter, indicate current frame id of run. */
|
||||
int32_t non_block; /* block flag of run, 0 is block else 1 is non block */
|
||||
int32_t timeout_ms; /* timeout for block mode, in milliseconds */
|
||||
int32_t fence_fd; /* fence fd from other unit */
|
||||
} rknn_run_extend;
|
||||
|
||||
/*
|
||||
the extend information for rknn_outputs_get.
|
||||
*/
|
||||
typedef struct _rknn_output_extend {
|
||||
uint64_t frame_id; /* output parameter, indicate the frame id of outputs, corresponds to
|
||||
struct rknn_run_extend.frame_id.*/
|
||||
} rknn_output_extend;
|
||||
|
||||
|
||||
/* rknn_init
|
||||
|
||||
initial the context and load the rknn model.
|
||||
|
||||
input:
|
||||
rknn_context* context the pointer of context handle.
|
||||
void* model if size > 0, pointer to the rknn model, if size = 0, filepath to the rknn model.
|
||||
uint32_t size the size of rknn model.
|
||||
uint32_t flag extend flag, see the define of RKNN_FLAG_XXX_XXX.
|
||||
rknn_init_extend* extend the extend information of init.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_init(rknn_context* context, void* model, uint32_t size, uint32_t flag, rknn_init_extend* extend);
|
||||
|
||||
/* rknn_dup_context
|
||||
|
||||
initial the context and load the rknn model.
|
||||
|
||||
input:
|
||||
rknn_context* context_in the pointer of context in handle.
|
||||
rknn_context* context_out the pointer of context out handle.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_dup_context(rknn_context* context_in, rknn_context* context_out);
|
||||
|
||||
/* rknn_destroy
|
||||
|
||||
unload the rknn model and destroy the context.
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_destroy(rknn_context context);
|
||||
|
||||
|
||||
/* rknn_query
|
||||
|
||||
query the information about model or others. see rknn_query_cmd.
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
rknn_query_cmd cmd the command of query.
|
||||
void* info the buffer point of information.
|
||||
uint32_t size the size of information.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_query(rknn_context context, rknn_query_cmd cmd, void* info, uint32_t size);
|
||||
|
||||
|
||||
/* rknn_inputs_set
|
||||
|
||||
set inputs information by input index of rknn model.
|
||||
inputs information see rknn_input.
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
uint32_t n_inputs the number of inputs.
|
||||
rknn_input inputs[] the arrays of inputs information, see rknn_input.
|
||||
return:
|
||||
int error code
|
||||
*/
|
||||
int rknn_inputs_set(rknn_context context, uint32_t n_inputs, rknn_input inputs[]);
|
||||
|
||||
/*
|
||||
rknn_set_batch_core_num
|
||||
|
||||
set rknn batch core_num.
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
int core_num the core number.
|
||||
return:
|
||||
int error code.
|
||||
|
||||
*/
|
||||
int rknn_set_batch_core_num(rknn_context context, int core_num);
|
||||
|
||||
/* rknn_set_core_mask
|
||||
|
||||
set the core mask for the model.(only supported on multi-core NPU platform)
|
||||
|
||||
RKNN_NPU_CORE_AUTO: auto mode, default value
|
||||
RKNN_NPU_CORE_0: core 0 mode
|
||||
RKNN_NPU_CORE_1: core 1 mode
|
||||
RKNN_NPU_CORE_2: core 2 mode
|
||||
RKNN_NPU_CORE_0_1: combine core 0/1 mode
|
||||
RKNN_NPU_CORE_0_1_2: combine core 0/1/2 mode
|
||||
RKNN_NPU_CORE_ALL: auto mode, select multiple npu cores to run depending on platform
|
||||
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
rknn_core_mask core_mask the core mask.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_set_core_mask(rknn_context context, rknn_core_mask core_mask);
|
||||
|
||||
/* rknn_run
|
||||
|
||||
run the model to execute inference.
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
rknn_run_extend* extend the extend information of run.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_run(rknn_context context, rknn_run_extend* extend);
|
||||
|
||||
|
||||
/* rknn_wait
|
||||
|
||||
wait the model after execute inference.
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
rknn_run_extend* extend the extend information of run.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_wait(rknn_context context, rknn_run_extend* extend);
|
||||
|
||||
|
||||
/* rknn_outputs_get
|
||||
|
||||
wait the inference to finish and get the outputs.
|
||||
this function will block until inference finish.
|
||||
the results will set to outputs[].
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
uint32_t n_outputs the number of outputs.
|
||||
rknn_output outputs[] the arrays of output, see rknn_output.
|
||||
rknn_output_extend* the extend information of output.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_outputs_get(rknn_context context, uint32_t n_outputs, rknn_output outputs[], rknn_output_extend* extend);
|
||||
|
||||
|
||||
/* rknn_outputs_release
|
||||
|
||||
release the outputs that get by rknn_outputs_get.
|
||||
after called, the rknn_output[x].buf get from rknn_outputs_get will
|
||||
also be free when rknn_output[x].is_prealloc = FALSE.
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
uint32_t n_ouputs the number of outputs.
|
||||
rknn_output outputs[] the arrays of output.
|
||||
return:
|
||||
int error code
|
||||
*/
|
||||
int rknn_outputs_release(rknn_context context, uint32_t n_ouputs, rknn_output outputs[]);
|
||||
|
||||
|
||||
/* new api for zero copy */
|
||||
|
||||
/* rknn_create_mem_from_phys (memory allocated outside)
|
||||
|
||||
initialize tensor memory from physical address.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
uint64_t phys_addr physical address.
|
||||
void *virt_addr virtual address.
|
||||
uint32_t size the size of tensor buffer.
|
||||
return:
|
||||
rknn_tensor_mem the pointer of tensor memory information.
|
||||
*/
|
||||
rknn_tensor_mem* rknn_create_mem_from_phys(rknn_context ctx, uint64_t phys_addr, void *virt_addr, uint32_t size);
|
||||
|
||||
|
||||
/* rknn_create_mem_from_fd (memory allocated outside)
|
||||
|
||||
initialize tensor memory from file description.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
int32_t fd file description.
|
||||
void *virt_addr virtual address.
|
||||
uint32_t size the size of tensor buffer.
|
||||
int32_t offset indicates the offset of the memory (virt_addr without offset).
|
||||
return:
|
||||
rknn_tensor_mem the pointer of tensor memory information.
|
||||
*/
|
||||
rknn_tensor_mem* rknn_create_mem_from_fd(rknn_context ctx, int32_t fd, void *virt_addr, uint32_t size, int32_t offset);
|
||||
|
||||
|
||||
/* rknn_create_mem_from_mb_blk (memory allocated outside)
|
||||
|
||||
create tensor memory from mb_blk.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
void *mb_blk mb_blk allocate from system api.
|
||||
int32_t offset indicates the offset of the memory.
|
||||
return:
|
||||
rknn_tensor_mem the pointer of tensor memory information.
|
||||
*/
|
||||
rknn_tensor_mem* rknn_create_mem_from_mb_blk(rknn_context ctx, void *mb_blk, int32_t offset);
|
||||
|
||||
|
||||
/* rknn_create_mem (memory allocated inside)
|
||||
|
||||
create tensor memory.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
uint32_t size the size of tensor buffer.
|
||||
return:
|
||||
rknn_tensor_mem the pointer of tensor memory information.
|
||||
*/
|
||||
rknn_tensor_mem* rknn_create_mem(rknn_context ctx, uint32_t size);
|
||||
|
||||
/* rknn_create_mem2 (memory allocated inside)
|
||||
|
||||
create tensor memory.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
uint64_t size the size of tensor buffer.
|
||||
uint64_t alloc_flags memory allocation flags.
|
||||
return:
|
||||
rknn_tensor_mem the pointer of tensor memory information.
|
||||
*/
|
||||
rknn_tensor_mem* rknn_create_mem2(rknn_context ctx, uint64_t size, uint64_t alloc_flags);
|
||||
|
||||
/* rknn_destroy_mem (support allocate inside and outside)
|
||||
|
||||
destroy tensor memory.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
rknn_tensor_mem *mem the pointer of tensor memory information.
|
||||
return:
|
||||
int error code
|
||||
*/
|
||||
int rknn_destroy_mem(rknn_context ctx, rknn_tensor_mem *mem);
|
||||
|
||||
|
||||
/* rknn_set_weight_mem
|
||||
|
||||
set the weight memory.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
rknn_tensor_mem *mem the array of tensor memory information
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_set_weight_mem(rknn_context ctx, rknn_tensor_mem *mem);
|
||||
|
||||
|
||||
/* rknn_set_internal_mem
|
||||
|
||||
set the internal memory.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
rknn_tensor_mem *mem the array of tensor memory information
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_set_internal_mem(rknn_context ctx, rknn_tensor_mem *mem);
|
||||
|
||||
|
||||
/* rknn_set_io_mem
|
||||
|
||||
set the input and output tensors buffer.
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
rknn_tensor_mem *mem the array of tensor memory information.
|
||||
rknn_tensor_attr *attr the attribute of input or output tensor buffer.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_set_io_mem(rknn_context ctx, rknn_tensor_mem *mem, rknn_tensor_attr *attr);
|
||||
|
||||
/* rknn_set_input_shape(deprecated)
|
||||
|
||||
set the input tensor shape (only valid for dynamic shape rknn model).
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
rknn_tensor_attr *attr the attribute of input or output tensor buffer.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_set_input_shape(rknn_context ctx, rknn_tensor_attr* attr);
|
||||
|
||||
/* rknn_set_input_shapes
|
||||
|
||||
set all the input tensor shapes. graph will run under current set of input shapes after rknn_set_input_shapes.(only valid for dynamic shape rknn model).
|
||||
|
||||
input:
|
||||
rknn_context ctx the handle of context.
|
||||
uint32_t n_inputs the number of inputs.
|
||||
rknn_tensor_attr attr[] the attribute array of all input tensors.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_set_input_shapes(rknn_context ctx, uint32_t n_inputs, rknn_tensor_attr attr[]);
|
||||
|
||||
/* rknn_mem_sync
|
||||
|
||||
sync cacheable rknn memory when both cpu and device access data.
|
||||
|
||||
input:
|
||||
rknn_context context the handle of context.
|
||||
rknn_tensor_mem *mem the pointer of tensor memory information.
|
||||
rknn_mem_sync_mode mode the mode of sync cache.
|
||||
return:
|
||||
int error code.
|
||||
*/
|
||||
int rknn_mem_sync(rknn_context context, rknn_tensor_mem* mem, rknn_mem_sync_mode mode);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} //extern "C"
|
||||
#endif
|
||||
|
||||
#endif //_RKNN_API_H
|
||||
7762
include/stb_image.h
Normal file
7762
include/stb_image.h
Normal file
File diff suppressed because it is too large
Load diff
4
models/SHA256SUMS
Normal file
4
models/SHA256SUMS
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79 models/onnx/face_recognition_sface_2021dec.onnx
|
||||
857efab2e0a5184ec86ffa7d0bf33ac94da92591e7650d5353622ce367218faf models/onnx/SCRFD_500M_KPS_640.onnx
|
||||
5f36840c6fea8a4772a45fe5eb3456b7bd2031e2986947154fcbd165f918f2ec models/rknn/face_recognition_sface_2021dec.rknn
|
||||
7d74abdedebc5fe25c98195db75cc915df74dd58d941be60c1c260186ac764e2 models/rknn/SCRFD_500M_KPS_640.rknn
|
||||
BIN
models/onnx/SCRFD_500M_KPS_640.onnx
Normal file
BIN
models/onnx/SCRFD_500M_KPS_640.onnx
Normal file
Binary file not shown.
BIN
models/onnx/face_recognition_sface_2021dec.onnx
Normal file
BIN
models/onnx/face_recognition_sface_2021dec.onnx
Normal file
Binary file not shown.
BIN
models/rknn/SCRFD_500M_KPS_640.rknn
Normal file
BIN
models/rknn/SCRFD_500M_KPS_640.rknn
Normal file
Binary file not shown.
BIN
models/rknn/face_recognition_sface_2021dec.rknn
Normal file
BIN
models/rknn/face_recognition_sface_2021dec.rknn
Normal file
Binary file not shown.
16
scripts/build.sh
Executable file
16
scripts/build.sh
Executable file
|
|
@ -0,0 +1,16 @@
|
|||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
|
||||
mkdir -p "$ROOT/bin"
|
||||
|
||||
g++ -O2 \
|
||||
"$ROOT/src/face_recognition.cc" \
|
||||
-I"$ROOT/include" \
|
||||
-L"$ROOT/runtime" \
|
||||
-Wl,-rpath,'$ORIGIN/../runtime' \
|
||||
-lrknnrt -ldl -lpthread \
|
||||
-o "$ROOT/bin/face_recognition"
|
||||
|
||||
echo "Built: $ROOT/bin/face_recognition"
|
||||
11
scripts/test.sh
Executable file
11
scripts/test.sh
Executable file
|
|
@ -0,0 +1,11 @@
|
|||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
|
||||
IMAGE1="${1:-$ROOT/test/test3f.jpg}"
|
||||
IMAGE2="${2:-$ROOT/test/test3f.jpg}"
|
||||
|
||||
cd "$ROOT"
|
||||
|
||||
"$ROOT/bin/face_recognition" "$IMAGE1" "$IMAGE2"
|
||||
786
src/face_recognition.cc
Normal file
786
src/face_recognition.cc
Normal file
|
|
@ -0,0 +1,786 @@
|
|||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#define STB_IMAGE_IMPLEMENTATION
|
||||
#include "stb_image.h"
|
||||
|
||||
#include "rknn_api.h"
|
||||
|
||||
static constexpr int SCRFD_SIZE = 640;
|
||||
static constexpr int SFACE_SIZE = 112;
|
||||
static constexpr int EMBED_DIM = 128;
|
||||
|
||||
struct Point {
|
||||
float x;
|
||||
float y;
|
||||
};
|
||||
|
||||
struct Face {
|
||||
float score;
|
||||
float x1, y1, x2, y2;
|
||||
Point kps[5];
|
||||
std::vector<float> embedding;
|
||||
};
|
||||
|
||||
static uint16_t float_to_fp16(float value)
|
||||
{
|
||||
uint32_t bits;
|
||||
std::memcpy(&bits, &value, sizeof(bits));
|
||||
|
||||
uint32_t sign = (bits >> 31) & 1;
|
||||
int exp = ((bits >> 23) & 0xff) - 127;
|
||||
uint32_t mant = bits & 0x7fffff;
|
||||
|
||||
if (exp == 128) {
|
||||
return (sign << 15) | 0x7c00;
|
||||
}
|
||||
|
||||
if (exp > 15) {
|
||||
return (sign << 15) | 0x7c00;
|
||||
}
|
||||
|
||||
if (exp < -14) {
|
||||
if (exp < -24)
|
||||
return sign << 15;
|
||||
|
||||
mant |= 0x800000;
|
||||
int shift = -exp - 14;
|
||||
uint16_t m = mant >> (shift + 13);
|
||||
|
||||
return (sign << 15) | m;
|
||||
}
|
||||
|
||||
uint16_t h_exp = (uint16_t)(exp + 15);
|
||||
uint16_t h_mant = (uint16_t)(mant >> 13);
|
||||
|
||||
return (sign << 15) | (h_exp << 10) | h_mant;
|
||||
}
|
||||
|
||||
static float fp16_to_float(uint16_t h)
|
||||
{
|
||||
uint32_t sign = (h >> 15) & 1;
|
||||
uint32_t exp = (h >> 10) & 0x1f;
|
||||
uint32_t mant = h & 0x3ff;
|
||||
|
||||
uint32_t bits;
|
||||
|
||||
if (exp == 0) {
|
||||
if (mant == 0) {
|
||||
bits = sign << 31;
|
||||
} else {
|
||||
float v = mant / 1024.0f;
|
||||
v = std::ldexp(v, -14);
|
||||
return sign ? -v : v;
|
||||
}
|
||||
} else if (exp == 31) {
|
||||
bits = (sign << 31) | 0x7f800000 | (mant << 13);
|
||||
} else {
|
||||
uint32_t fexp = exp - 15 + 127;
|
||||
bits = (sign << 31) | (fexp << 23) | (mant << 13);
|
||||
}
|
||||
|
||||
float v;
|
||||
std::memcpy(&v, &bits, sizeof(v));
|
||||
return v;
|
||||
}
|
||||
|
||||
static float iou(const Face& a, const Face& b)
|
||||
{
|
||||
float xx1 = std::max(a.x1, b.x1);
|
||||
float yy1 = std::max(a.y1, b.y1);
|
||||
float xx2 = std::min(a.x2, b.x2);
|
||||
float yy2 = std::min(a.y2, b.y2);
|
||||
|
||||
float w = std::max(0.0f, xx2 - xx1);
|
||||
float h = std::max(0.0f, yy2 - yy1);
|
||||
float inter = w * h;
|
||||
|
||||
float area_a = std::max(0.0f, a.x2-a.x1) *
|
||||
std::max(0.0f, a.y2-a.y1);
|
||||
|
||||
float area_b = std::max(0.0f, b.x2-b.x1) *
|
||||
std::max(0.0f, b.y2-b.y1);
|
||||
|
||||
return inter / (area_a + area_b - inter + 1e-6f);
|
||||
}
|
||||
|
||||
static void nms(std::vector<Face>& faces, float threshold)
|
||||
{
|
||||
std::sort(faces.begin(), faces.end(),
|
||||
[](const Face& a, const Face& b) {
|
||||
return a.score > b.score;
|
||||
});
|
||||
|
||||
std::vector<Face> result;
|
||||
|
||||
for (const auto& f : faces) {
|
||||
bool keep = true;
|
||||
|
||||
for (const auto& r : result) {
|
||||
if (iou(f, r) > threshold) {
|
||||
keep = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (keep)
|
||||
result.push_back(f);
|
||||
}
|
||||
|
||||
faces.swap(result);
|
||||
}
|
||||
|
||||
struct Image {
|
||||
int w = 0;
|
||||
int h = 0;
|
||||
std::vector<uint8_t> rgb;
|
||||
};
|
||||
|
||||
static bool load_image(const std::string& path, Image& img)
|
||||
{
|
||||
int c = 0;
|
||||
|
||||
unsigned char* p =
|
||||
stbi_load(path.c_str(), &img.w, &img.h, &c, 3);
|
||||
|
||||
if (!p) {
|
||||
std::cerr << "Errore caricamento: " << path << "\n";
|
||||
return false;
|
||||
}
|
||||
|
||||
img.rgb.assign(p, p + img.w * img.h * 3);
|
||||
stbi_image_free(p);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
static void resize_bilinear_rgb(
|
||||
const Image& src,
|
||||
std::vector<uint8_t>& dst,
|
||||
int dw,
|
||||
int dh)
|
||||
{
|
||||
dst.resize(dw * dh * 3);
|
||||
|
||||
float sx = (float)src.w / dw;
|
||||
float sy = (float)src.h / dh;
|
||||
|
||||
for (int y = 0; y < dh; ++y) {
|
||||
float fy = (y + 0.5f) * sy - 0.5f;
|
||||
int y0 = (int)std::floor(fy);
|
||||
float wy = fy - y0;
|
||||
|
||||
if (y0 < 0) {
|
||||
y0 = 0;
|
||||
wy = 0;
|
||||
}
|
||||
|
||||
int y1 = std::min(y0 + 1, src.h - 1);
|
||||
|
||||
for (int x = 0; x < dw; ++x) {
|
||||
float fx = (x + 0.5f) * sx - 0.5f;
|
||||
int x0 = (int)std::floor(fx);
|
||||
float wx = fx - x0;
|
||||
|
||||
if (x0 < 0) {
|
||||
x0 = 0;
|
||||
wx = 0;
|
||||
}
|
||||
|
||||
int x1 = std::min(x0 + 1, src.w - 1);
|
||||
|
||||
for (int c = 0; c < 3; ++c) {
|
||||
float p00 = src.rgb[(y0*src.w+x0)*3+c];
|
||||
float p01 = src.rgb[(y0*src.w+x1)*3+c];
|
||||
float p10 = src.rgb[(y1*src.w+x0)*3+c];
|
||||
float p11 = src.rgb[(y1*src.w+x1)*3+c];
|
||||
|
||||
float v =
|
||||
p00 * (1-wx) * (1-wy) +
|
||||
p01 * wx * (1-wy) +
|
||||
p10 * (1-wx) * wy +
|
||||
p11 * wx * wy;
|
||||
|
||||
dst[(y*dw+x)*3+c] =
|
||||
(uint8_t)std::clamp((int)std::round(v), 0, 255);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
class Scrfd {
|
||||
public:
|
||||
rknn_context ctx = 0;
|
||||
|
||||
bool init(const char* model)
|
||||
{
|
||||
FILE* fp = fopen(model, "rb");
|
||||
if (!fp) {
|
||||
perror(model);
|
||||
return false;
|
||||
}
|
||||
|
||||
fseek(fp, 0, SEEK_END);
|
||||
size_t size = ftell(fp);
|
||||
fseek(fp, 0, SEEK_SET);
|
||||
|
||||
std::vector<uint8_t> data(size);
|
||||
fread(data.data(), 1, size, fp);
|
||||
fclose(fp);
|
||||
|
||||
int ret = rknn_init(&ctx, data.data(), size, 0, nullptr);
|
||||
if (ret != 0) {
|
||||
std::cerr << "rknn_init SCRFD failed: " << ret << "\n";
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<Face> detect(const Image& image)
|
||||
{
|
||||
float scale =
|
||||
std::min(640.0f / image.w, 640.0f / image.h);
|
||||
|
||||
int rw = (int)std::round(image.w * scale);
|
||||
int rh = (int)std::round(image.h * scale);
|
||||
|
||||
std::vector<uint8_t> resized;
|
||||
resize_bilinear_rgb(image, resized, rw, rh);
|
||||
|
||||
std::vector<uint16_t> input(
|
||||
SCRFD_SIZE * SCRFD_SIZE * 3);
|
||||
|
||||
for (int y = 0; y < rh; ++y) {
|
||||
for (int x = 0; x < rw; ++x) {
|
||||
int dst = (y * SCRFD_SIZE + x) * 3;
|
||||
int src = (y * rw + x) * 3;
|
||||
|
||||
for (int c = 0; c < 3; ++c) {
|
||||
float v =
|
||||
((float)resized[src+c] - 127.5f) / 128.0f;
|
||||
|
||||
input[dst+c] = float_to_fp16(v);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
rknn_input in{};
|
||||
in.index = 0;
|
||||
in.type = RKNN_TENSOR_FLOAT16;
|
||||
in.fmt = RKNN_TENSOR_NHWC;
|
||||
in.size = input.size() * sizeof(uint16_t);
|
||||
in.buf = input.data();
|
||||
in.pass_through = 1;
|
||||
|
||||
int ret = rknn_inputs_set(ctx, 1, &in);
|
||||
if (ret != 0) {
|
||||
std::cerr << "SCRFD inputs_set failed\n";
|
||||
return {};
|
||||
}
|
||||
|
||||
ret = rknn_run(ctx, nullptr);
|
||||
if (ret != 0) {
|
||||
std::cerr << "SCRFD run failed\n";
|
||||
return {};
|
||||
}
|
||||
|
||||
rknn_output outputs[9]{};
|
||||
|
||||
for (int i = 0; i < 9; ++i) {
|
||||
outputs[i].want_float = 0;
|
||||
}
|
||||
|
||||
ret = rknn_outputs_get(ctx, 9, outputs, nullptr);
|
||||
if (ret != 0) {
|
||||
std::cerr << "SCRFD outputs_get failed\n";
|
||||
return {};
|
||||
}
|
||||
|
||||
const int strides[3] = {8, 16, 32};
|
||||
const int counts[3] = {12800, 3200, 800};
|
||||
|
||||
std::vector<Face> faces;
|
||||
|
||||
const float score_threshold = 0.50f;
|
||||
|
||||
for (int level = 0; level < 3; ++level) {
|
||||
const int stride = strides[level];
|
||||
const int count = counts[level];
|
||||
|
||||
const uint16_t* scores =
|
||||
(const uint16_t*)outputs[level].buf;
|
||||
|
||||
const uint16_t* bbox =
|
||||
(const uint16_t*)outputs[3 + level].buf;
|
||||
|
||||
const uint16_t* kps =
|
||||
(const uint16_t*)outputs[6 + level].buf;
|
||||
|
||||
int feat_w = SCRFD_SIZE / stride;
|
||||
int feat_h = SCRFD_SIZE / stride;
|
||||
|
||||
for (int i = 0; i < count; ++i) {
|
||||
float score = fp16_to_float(scores[i]);
|
||||
|
||||
if (score < score_threshold)
|
||||
continue;
|
||||
|
||||
int anchor = i % 2;
|
||||
int p = i / 2;
|
||||
|
||||
int gx = p % feat_w;
|
||||
int gy = p / feat_w;
|
||||
|
||||
float cx = gx * stride;
|
||||
float cy = gy * stride;
|
||||
|
||||
/*
|
||||
* AnchorGenerator con due anchor:
|
||||
* centers duplicati sullo stesso punto.
|
||||
*/
|
||||
|
||||
(void)anchor;
|
||||
|
||||
float l = fp16_to_float(bbox[i*4+0]) * stride;
|
||||
float t = fp16_to_float(bbox[i*4+1]) * stride;
|
||||
float r = fp16_to_float(bbox[i*4+2]) * stride;
|
||||
float b = fp16_to_float(bbox[i*4+3]) * stride;
|
||||
|
||||
Face f;
|
||||
f.score = score;
|
||||
|
||||
f.x1 = (cx - l) / scale;
|
||||
f.y1 = (cy - t) / scale;
|
||||
f.x2 = (cx + r) / scale;
|
||||
f.y2 = (cy + b) / scale;
|
||||
|
||||
for (int j = 0; j < 5; ++j) {
|
||||
float x =
|
||||
cx + fp16_to_float(kps[i*10+j*2])
|
||||
* stride;
|
||||
|
||||
float y =
|
||||
cy + fp16_to_float(kps[i*10+j*2+1])
|
||||
* stride;
|
||||
|
||||
f.kps[j].x = x / scale;
|
||||
f.kps[j].y = y / scale;
|
||||
}
|
||||
|
||||
f.x1 = std::clamp(f.x1, 0.0f, (float)image.w);
|
||||
f.y1 = std::clamp(f.y1, 0.0f, (float)image.h);
|
||||
f.x2 = std::clamp(f.x2, 0.0f, (float)image.w);
|
||||
f.y2 = std::clamp(f.y2, 0.0f, (float)image.h);
|
||||
|
||||
faces.push_back(f);
|
||||
}
|
||||
}
|
||||
|
||||
rknn_outputs_release(ctx, 9, outputs);
|
||||
|
||||
nms(faces, 0.45f);
|
||||
|
||||
return faces;
|
||||
}
|
||||
|
||||
~Scrfd()
|
||||
{
|
||||
if (ctx)
|
||||
rknn_destroy(ctx);
|
||||
}
|
||||
};
|
||||
|
||||
static const Point sface_template[5] = {
|
||||
{38.2946f, 51.6963f},
|
||||
{73.5318f, 51.5014f},
|
||||
{56.0252f, 71.7366f},
|
||||
{41.5493f, 92.3655f},
|
||||
{70.7299f, 92.2041f}
|
||||
};
|
||||
|
||||
/*
|
||||
* Trova la trasformazione similarity:
|
||||
*
|
||||
* x' = a*x - b*y + tx
|
||||
* y' = b*x + a*y + ty
|
||||
*/
|
||||
static void similarity_transform(
|
||||
const Point src[5],
|
||||
float& a,
|
||||
float& b,
|
||||
float& tx,
|
||||
float& ty)
|
||||
{
|
||||
double sx = 0, sy = 0;
|
||||
double dx = 0, dy = 0;
|
||||
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
sx += src[i].x;
|
||||
sy += src[i].y;
|
||||
dx += sface_template[i].x;
|
||||
dy += sface_template[i].y;
|
||||
}
|
||||
|
||||
sx /= 5;
|
||||
sy /= 5;
|
||||
dx /= 5;
|
||||
dy /= 5;
|
||||
|
||||
double num_a = 0;
|
||||
double num_b = 0;
|
||||
double den = 0;
|
||||
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
double x = src[i].x - sx;
|
||||
double y = src[i].y - sy;
|
||||
|
||||
double X = sface_template[i].x - dx;
|
||||
double Y = sface_template[i].y - dy;
|
||||
|
||||
num_a += x*X + y*Y;
|
||||
num_b += x*Y - y*X;
|
||||
den += x*x + y*y;
|
||||
}
|
||||
|
||||
a = (float)(num_a / den);
|
||||
b = (float)(num_b / den);
|
||||
|
||||
tx = (float)(dx - a*sx + b*sy);
|
||||
ty = (float)(dy - b*sx - a*sy);
|
||||
}
|
||||
|
||||
static std::vector<uint16_t> align_face(
|
||||
const Image& image,
|
||||
const Face& face)
|
||||
{
|
||||
float a, b, tx, ty;
|
||||
|
||||
similarity_transform(face.kps, a, b, tx, ty);
|
||||
|
||||
/*
|
||||
* Inversa della similarity transform.
|
||||
*
|
||||
* src = A^-1 * (dst - t)
|
||||
*/
|
||||
float denom = a*a + b*b;
|
||||
|
||||
std::vector<uint16_t> output(
|
||||
SFACE_SIZE * SFACE_SIZE * 3);
|
||||
|
||||
for (int y = 0; y < SFACE_SIZE; ++y) {
|
||||
for (int x = 0; x < SFACE_SIZE; ++x) {
|
||||
|
||||
float X = x - tx;
|
||||
float Y = y - ty;
|
||||
|
||||
float sx =
|
||||
(a*X + b*Y) / denom;
|
||||
|
||||
float sy =
|
||||
(-b*X + a*Y) / denom;
|
||||
|
||||
int x0 = (int)std::floor(sx);
|
||||
int y0 = (int)std::floor(sy);
|
||||
|
||||
float wx = sx - x0;
|
||||
float wy = sy - y0;
|
||||
|
||||
uint8_t rgb[3] = {0,0,0};
|
||||
|
||||
if (x0 >= 0 &&
|
||||
y0 >= 0 &&
|
||||
x0 + 1 < image.w &&
|
||||
y0 + 1 < image.h) {
|
||||
|
||||
int x1 = x0 + 1;
|
||||
int y1 = y0 + 1;
|
||||
|
||||
for (int c = 0; c < 3; ++c) {
|
||||
float p00 =
|
||||
image.rgb[(y0*image.w+x0)*3+c];
|
||||
|
||||
float p01 =
|
||||
image.rgb[(y0*image.w+x1)*3+c];
|
||||
|
||||
float p10 =
|
||||
image.rgb[(y1*image.w+x0)*3+c];
|
||||
|
||||
float p11 =
|
||||
image.rgb[(y1*image.w+x1)*3+c];
|
||||
|
||||
float v =
|
||||
p00*(1-wx)*(1-wy) +
|
||||
p01*wx*(1-wy) +
|
||||
p10*(1-wx)*wy +
|
||||
p11*wx*wy;
|
||||
|
||||
rgb[c] =
|
||||
(uint8_t)std::clamp(
|
||||
(int)std::round(v), 0, 255);
|
||||
}
|
||||
}
|
||||
|
||||
int idx = (y*SFACE_SIZE+x)*3;
|
||||
|
||||
/*
|
||||
* SFace OpenCV:
|
||||
* blobFromImage(..., swapRB=true, scalefactor=1)
|
||||
*
|
||||
* Noi partiamo già da RGB, quindi non facciamo
|
||||
* nessun ulteriore swap e nessuna normalizzazione.
|
||||
*/
|
||||
output[idx+0] = float_to_fp16(rgb[0]);
|
||||
output[idx+1] = float_to_fp16(rgb[1]);
|
||||
output[idx+2] = float_to_fp16(rgb[2]);
|
||||
}
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
class SFace {
|
||||
public:
|
||||
rknn_context ctx = 0;
|
||||
|
||||
bool init(const char* model)
|
||||
{
|
||||
FILE* fp = fopen(model, "rb");
|
||||
if (!fp) {
|
||||
perror(model);
|
||||
return false;
|
||||
}
|
||||
|
||||
fseek(fp, 0, SEEK_END);
|
||||
size_t size = ftell(fp);
|
||||
fseek(fp, 0, SEEK_SET);
|
||||
|
||||
std::vector<uint8_t> data(size);
|
||||
fread(data.data(), 1, size, fp);
|
||||
fclose(fp);
|
||||
|
||||
int ret = rknn_init(&ctx, data.data(), size, 0, nullptr);
|
||||
|
||||
if (ret != 0) {
|
||||
std::cerr << "rknn_init SFace failed: "
|
||||
<< ret << "\n";
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<float> feature(
|
||||
const std::vector<uint16_t>& input)
|
||||
{
|
||||
rknn_input in{};
|
||||
|
||||
in.index = 0;
|
||||
in.type = RKNN_TENSOR_FLOAT16;
|
||||
in.fmt = RKNN_TENSOR_NHWC;
|
||||
in.size = input.size() * sizeof(uint16_t);
|
||||
in.buf = (void*)input.data();
|
||||
in.pass_through = 1;
|
||||
|
||||
int ret = rknn_inputs_set(ctx, 1, &in);
|
||||
|
||||
if (ret != 0) {
|
||||
std::cerr << "SFace inputs_set failed\n";
|
||||
return {};
|
||||
}
|
||||
|
||||
ret = rknn_run(ctx, nullptr);
|
||||
|
||||
if (ret != 0) {
|
||||
std::cerr << "SFace run failed\n";
|
||||
return {};
|
||||
}
|
||||
|
||||
rknn_output output{};
|
||||
|
||||
output.want_float = 0;
|
||||
|
||||
ret = rknn_outputs_get(
|
||||
ctx, 1, &output, nullptr);
|
||||
|
||||
if (ret != 0) {
|
||||
std::cerr << "SFace outputs_get failed\n";
|
||||
return {};
|
||||
}
|
||||
|
||||
const uint16_t* p =
|
||||
(const uint16_t*)output.buf;
|
||||
|
||||
std::vector<float> emb(EMBED_DIM);
|
||||
|
||||
for (int i = 0; i < EMBED_DIM; ++i)
|
||||
emb[i] = fp16_to_float(p[i]);
|
||||
|
||||
rknn_outputs_release(ctx, 1, &output);
|
||||
|
||||
/*
|
||||
* L2 normalization.
|
||||
*/
|
||||
double norm = 0;
|
||||
|
||||
for (float v : emb)
|
||||
norm += (double)v * v;
|
||||
|
||||
norm = std::sqrt(norm);
|
||||
|
||||
if (norm > 0) {
|
||||
for (float& v : emb)
|
||||
v /= (float)norm;
|
||||
}
|
||||
|
||||
return emb;
|
||||
}
|
||||
|
||||
~SFace()
|
||||
{
|
||||
if (ctx)
|
||||
rknn_destroy(ctx);
|
||||
}
|
||||
};
|
||||
|
||||
static float cosine(
|
||||
const std::vector<float>& a,
|
||||
const std::vector<float>& b)
|
||||
{
|
||||
float s = 0;
|
||||
|
||||
for (size_t i = 0; i < a.size(); ++i)
|
||||
s += a[i] * b[i];
|
||||
|
||||
return s;
|
||||
}
|
||||
|
||||
static std::vector<Face> process_image(
|
||||
const Image& image,
|
||||
Scrfd& detector,
|
||||
SFace& recognizer)
|
||||
{
|
||||
auto faces = detector.detect(image);
|
||||
|
||||
for (auto& face : faces) {
|
||||
auto aligned = align_face(image, face);
|
||||
face.embedding =
|
||||
recognizer.feature(aligned);
|
||||
}
|
||||
|
||||
return faces;
|
||||
}
|
||||
|
||||
static void print_faces(
|
||||
const std::string& name,
|
||||
const std::vector<Face>& faces)
|
||||
{
|
||||
std::cout << "\n" << name << "\n";
|
||||
std::cout << "Faces: " << faces.size() << "\n";
|
||||
|
||||
for (size_t i = 0; i < faces.size(); ++i) {
|
||||
const auto& f = faces[i];
|
||||
|
||||
std::cout
|
||||
<< "\nFace " << i
|
||||
<< " score=" << f.score
|
||||
<< "\n bbox="
|
||||
<< f.x1 << " "
|
||||
<< f.y1 << " "
|
||||
<< f.x2 << " "
|
||||
<< f.y2 << "\n";
|
||||
|
||||
std::cout << " landmarks:\n";
|
||||
|
||||
for (int j = 0; j < 5; ++j) {
|
||||
std::cout
|
||||
<< " " << j << ": "
|
||||
<< f.kps[j].x << " "
|
||||
<< f.kps[j].y << "\n";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
if (argc != 2 && argc != 3) {
|
||||
std::cerr
|
||||
<< "Uso:\n"
|
||||
<< " " << argv[0]
|
||||
<< " image.jpg\n"
|
||||
<< " " << argv[0]
|
||||
<< " reference.jpg query.jpg\n";
|
||||
return 1;
|
||||
}
|
||||
|
||||
const char* SCRFD_MODEL =
|
||||
"models/rknn/SCRFD_500M_KPS_640.rknn";
|
||||
|
||||
const char* SFACE_MODEL =
|
||||
"models/rknn/face_recognition_sface_2021dec.rknn";
|
||||
|
||||
Scrfd detector;
|
||||
|
||||
if (!detector.init(SCRFD_MODEL))
|
||||
return 1;
|
||||
|
||||
SFace recognizer;
|
||||
|
||||
if (!recognizer.init(SFACE_MODEL))
|
||||
return 1;
|
||||
|
||||
Image image1;
|
||||
|
||||
if (!load_image(argv[1], image1))
|
||||
return 1;
|
||||
|
||||
auto faces1 =
|
||||
process_image(image1, detector, recognizer);
|
||||
|
||||
print_faces(argv[1], faces1);
|
||||
|
||||
if (argc == 2)
|
||||
return 0;
|
||||
|
||||
Image image2;
|
||||
|
||||
if (!load_image(argv[2], image2))
|
||||
return 1;
|
||||
|
||||
auto faces2 =
|
||||
process_image(image2, detector, recognizer);
|
||||
|
||||
print_faces(argv[2], faces2);
|
||||
|
||||
std::cout << "\n========================================\n";
|
||||
std::cout << "COSINE SIMILARITY\n";
|
||||
std::cout << "========================================\n";
|
||||
|
||||
for (size_t i = 0; i < faces1.size(); ++i) {
|
||||
for (size_t j = 0; j < faces2.size(); ++j) {
|
||||
|
||||
float sim =
|
||||
cosine(
|
||||
faces1[i].embedding,
|
||||
faces2[j].embedding);
|
||||
|
||||
std::cout
|
||||
<< "ref[" << i << "] vs query[" << j << "]"
|
||||
<< " = " << sim;
|
||||
|
||||
if (sim >= 0.363f)
|
||||
std::cout << " MATCH";
|
||||
|
||||
std::cout << "\n";
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
BIN
test/test3f.jpg
Normal file
BIN
test/test3f.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 401 KiB |
73
tools/conversion/convert_legacy.py
Normal file
73
tools/conversion/convert_legacy.py
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
import sys
|
||||
from rknn.api import RKNN
|
||||
|
||||
DATASET_PATH = 'model/dataset.txt'
|
||||
DEFAULT_RKNN_PATH = '../model/RetinaFace.rknn'
|
||||
DEFAULT_QUANT = True
|
||||
|
||||
def parse_arg():
|
||||
if len(sys.argv) < 3:
|
||||
print("Usage: python3 {} onnx_model_path [platform] [dtype(optional)] [output_rknn_path(optional)]".format(sys.argv[0]));
|
||||
print(" platform choose from [rk3562, rk3566, rk3568, rk3576, rk3588, rv1126b, rv1109, rv1126, rk1808]")
|
||||
print(" dtype choose from [i8] for [rk3562, rk3566, rk3568, rk3576, rk3588, rv1126b]")
|
||||
print(" dtype choose from [u8] for [rv1109, rv1126, rk1808]")
|
||||
exit(1)
|
||||
|
||||
model_path = sys.argv[1]
|
||||
platform = sys.argv[2]
|
||||
|
||||
do_quant = DEFAULT_QUANT
|
||||
if len(sys.argv) > 3:
|
||||
model_type = sys.argv[3]
|
||||
if model_type not in ['i8', 'u8', 'fp']:
|
||||
print("ERROR: Invalid model type: {}".format(model_type))
|
||||
exit(1)
|
||||
elif model_type in ['i8', 'u8']:
|
||||
do_quant = True
|
||||
else:
|
||||
do_quant = False
|
||||
|
||||
if len(sys.argv) > 4:
|
||||
output_path = sys.argv[4]
|
||||
else:
|
||||
output_path = DEFAULT_RKNN_PATH
|
||||
|
||||
return model_path, platform, do_quant, output_path
|
||||
|
||||
if __name__ == '__main__':
|
||||
model_path, platform, do_quant, output_path = parse_arg()
|
||||
|
||||
# Create RKNN object
|
||||
rknn = RKNN(verbose=False)
|
||||
|
||||
# Pre-process config
|
||||
print('--> Config model')
|
||||
rknn.config(mean_values=[[104, 117, 123]], std_values=[[1, 1, 1]], target_platform=platform)
|
||||
print('done')
|
||||
|
||||
# Load model
|
||||
print('--> Loading model')
|
||||
ret = rknn.load_onnx(model=model_path)
|
||||
if ret != 0:
|
||||
print('Load model failed!')
|
||||
exit(ret)
|
||||
print('done')
|
||||
|
||||
# Build model
|
||||
print('--> Building model')
|
||||
ret = rknn.build(do_quantization=do_quant, dataset=DATASET_PATH)
|
||||
if ret != 0:
|
||||
print('Build model failed!')
|
||||
exit(ret)
|
||||
print('done')
|
||||
|
||||
# Export rknn model
|
||||
print('--> Export rknn model')
|
||||
ret = rknn.export_rknn(output_path)
|
||||
if ret != 0:
|
||||
print('Export rknn model failed!')
|
||||
exit(ret)
|
||||
print('done')
|
||||
|
||||
# Release
|
||||
rknn.release()
|
||||
45
tools/conversion/convert_scrfd_rknn.py
Normal file
45
tools/conversion/convert_scrfd_rknn.py
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
from rknn.api import RKNN
|
||||
|
||||
ONNX_MODEL = "SCRFD_500M_KPS_640.onnx"
|
||||
RKNN_MODEL = "SCRFD_500M_KPS_640.rknn"
|
||||
|
||||
print("=" * 60)
|
||||
print("SCRFD ONNX -> RKNN")
|
||||
print("=" * 60)
|
||||
|
||||
rknn = RKNN(verbose=True)
|
||||
|
||||
print("\n[1] Configuring RKNN...")
|
||||
ret = rknn.config(
|
||||
target_platform="rk3588"
|
||||
)
|
||||
|
||||
if ret != 0:
|
||||
raise RuntimeError("rknn.config() failed")
|
||||
|
||||
print("\n[2] Loading ONNX...")
|
||||
ret = rknn.load_onnx(
|
||||
model=ONNX_MODEL
|
||||
)
|
||||
|
||||
if ret != 0:
|
||||
raise RuntimeError("rknn.load_onnx() failed")
|
||||
|
||||
print("\n[3] Building RKNN (FP16, no quantization)...")
|
||||
ret = rknn.build(
|
||||
do_quantization=False
|
||||
)
|
||||
|
||||
if ret != 0:
|
||||
raise RuntimeError("rknn.build() failed")
|
||||
|
||||
print("\n[4] Exporting RKNN...")
|
||||
ret = rknn.export_rknn(RKNN_MODEL)
|
||||
|
||||
if ret != 0:
|
||||
raise RuntimeError("rknn.export_rknn() failed")
|
||||
|
||||
print("\nSUCCESS")
|
||||
print(RKNN_MODEL)
|
||||
|
||||
rknn.release()
|
||||
Loading…
Reference in a new issue