├── .gitignore ├── .gitmodules ├── LICENSE ├── README.md ├── model ├── bus.jpg ├── dataset.txt ├── docker │ └── dockerfile └── onnx2rknn.py ├── rknpu2_ros ├── CMakeLists.txt ├── README.md └── package.xml ├── rknpu2_ros_common ├── CMakeLists.txt ├── README.md ├── include │ └── rknpu2_ros_common │ │ ├── coco_names.hpp │ │ ├── drm_func.h │ │ ├── model_loader.h │ │ ├── object.hpp │ │ ├── rga_func.h │ │ └── utils.hpp └── package.xml └── rknpu2_ros_yolov5 ├── CMakeLists.txt ├── README.md ├── include └── rknpu2_ros_yolov5 │ ├── postprocess.hpp │ ├── rknpu2_yolov5.hpp │ └── rknpu2_yolov5_node.hpp ├── package.xml └── src ├── main.cc ├── postprocess.cc ├── postprocess.h ├── rknpu2_yolov5.cpp └── rknpu2_yolov5_node.cpp /.gitignore: -------------------------------------------------------------------------------- 1 | # Prerequisites 2 | *.d 3 | 4 | # Compiled Object files 5 | *.slo 6 | *.lo 7 | *.o 8 | *.obj 9 | 10 | # Precompiled Headers 11 | *.gch 12 | *.pch 13 | 14 | # Compiled Dynamic libraries 15 | *.so 16 | *.dylib 17 | *.dll 18 | 19 | # Fortran module files 20 | *.mod 21 | *.smod 22 | 23 | # Compiled Static libraries 24 | *.lai 25 | *.la 26 | *.a 27 | *.lib 28 | 29 | # Executables 30 | *.exe 31 | *.out 32 | *.app 33 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "rknpu2"] 2 | path = rknpu2 3 | url = https://github.com/rockchip-linux/rknpu2 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2022, fateshelled 4 | 5 | Redistribution and use in source and binary forms, with or without 6 | modification, are permitted provided that the following conditions are met: 7 | 8 | 1. Redistributions of source code must retain the above copyright notice, this 9 | list of conditions and the following disclaimer. 10 | 11 | 2. Redistributions in binary form must reproduce the above copyright notice, 12 | this list of conditions and the following disclaimer in the documentation 13 | and/or other materials provided with the distribution. 14 | 15 | 3. Neither the name of the copyright holder nor the names of its 16 | contributors may be used to endorse or promote products derived from 17 | this software without specific prior written permission. 18 | 19 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 20 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 21 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 22 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 23 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 24 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 25 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 26 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 27 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 29 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # rknpu2_ros 2 | ROS2 Inference sample for using Rockchip NPU. 3 | 4 | tested OrangePi5 (RK3588s) Ubuntu 22.04 + ROS2 Humble. 5 | 6 | ## build 7 | ```bash 8 | # clone repository 9 | cd ros2_ws/src 10 | git clone --recursive https://github.com/fateshelled/rknpu2_ros 11 | cd ../ 12 | 13 | # build 14 | # TARGET_SOC = RK3588 or RK356X or RV110X 15 | colcon build --symlink-install --packages-up-to rknpu2_ros --cmake-args -D TARGET_SOC=RK3588 16 | ``` 17 | 18 | ## run 19 | ```bash 20 | ros2 run v4l2_camera v4l2_camera_node 21 | 22 | ros2 run rknpu2_ros_yolov5 rknpu2_ros_yolov5 23 | ``` 24 | -------------------------------------------------------------------------------- /model/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fateshelled/rknpu2_ros/c44c548c20a5b0b85bd07e2e0b8f4c47c702a5b3/model/bus.jpg -------------------------------------------------------------------------------- /model/dataset.txt: -------------------------------------------------------------------------------- 1 | bus.jpg 2 | -------------------------------------------------------------------------------- /model/docker/dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:20.04 2 | ENV DEBIAN_FRONTEND=noninteractive 3 | 4 | RUN apt update && apt upgrade && \ 5 | apt install -y wget git python3-dev python3-pip && \ 6 | apt -y clean && \ 7 | rm -rf /var/lib/apt/lists/* 8 | 9 | WORKDIR /workspace 10 | 11 | RUN wget https://github.com/rockchip-linux/rknn-toolkit2/raw/master/packages/rknn_toolkit2-1.4.0_22dcfef4-cp38-cp38-linux_x86_64.whl &&\ 12 | python3 -m pip install -U pip && \ 13 | python3 -m pip install numpy==1.19.5 && \ 14 | python3 -m pip install rknn_toolkit2-1.4.0_22dcfef4-cp38-cp38-linux_x86_64.whl && \ 15 | python3 -m pip install opencv-python-headless && \ 16 | rm rknn_toolkit2-1.4.0_22dcfef4-cp38-cp38-linux_x86_64.whl && \ 17 | python3 -m pip cache purge 18 | -------------------------------------------------------------------------------- /model/onnx2rknn.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import argparse 4 | import ast 5 | 6 | from rknn.api import RKNN 7 | 8 | if __name__ == '__main__': 9 | 10 | parser = argparse.ArgumentParser(description='') 11 | parser.add_argument('--onnx_model_path', type=str, required=True, help='onnx model path') 12 | parser.add_argument('--platform', type=str, default="rk3588", help='platform name. Currently support rk3566 / rk3568 / rk3588 / rv1103 / rv1106. default="rk3588"') 13 | parser.add_argument('--output_dir', type=str, default="rknn_models", help='rknn model output directory. default="rknn_models"') 14 | parser.add_argument('--dataset_file', type=str, default="./dataset.txt", help='dataset file path. default="./dataset.txt"') 15 | args = parser.parse_args() 16 | 17 | PLATFORM = args.platform 18 | ONNX_MODEL_PATH = './yolox_s.onnx' 19 | EXP = os.path.splitext(os.path.basename(ONNX_MODEL_PATH))[0] 20 | OUT_DIR = args.output_dir 21 | DATASET = args.dataset_file 22 | RKNN_MODEL_PATH = './{}/{}.rknn'.format(OUT_DIR, EXP) 23 | # OPT_ONNX_MODEL_PATH = './{}/{}_opt.onnx'.format(OUT_DIR, EXP) 24 | MEAN_VALUES = [[0, 0, 0]] 25 | STD_VALUES = [[255, 255, 255]] 26 | 27 | print() 28 | print('--- Convert onnx model to rknn model ---') 29 | 30 | # Create RKNN object 31 | rknn = RKNN(verbose=False) 32 | 33 | # :param mean_values: Channel mean value list. 34 | # :param std_values: Channel std value list. 35 | # :param quantized_dtype: quantize data type, currently support: asymmetric_quantized-8. 36 | # :param quantized_algorithm: currently support: normal, mmse (Min Mean Square Error), kl_divergence. 37 | # :param quantized_method: quantize method, currently support: layer, channel. 38 | # :param target_platform: target chip platform, default is None, means target platform is rk3566. Currently support rk3566 / rk3568 / rk3588 / rv1103 / rv1106. 39 | # :param quant_img_RGB2BGR: whether to do RGB2BGR when load quantize image (jpg/jpeg/png/bmp), default is False. 40 | # :param float_dtype: non quantize data type, currently support: float16, default is float16. 41 | # :param optimization_level: set optimization level, default 3 means use all default optimization options. 42 | # :param custom_string: add custom string information to rknn model, then can query the information at runtime. 43 | # :param remove_weight: generate a slave rknn model which removes conv2d weight, need share weight with rknn model of complete weights. 44 | # :param compress_weight: compress the weights of the model, which can reduce the size of rknn model. 45 | # :param inputs_yuv_fmt: add yuv preprocess at the top of model. 46 | # :param single_core_mode: only for rk3588. single_core_mode=True can reduce the size of rknn model. 47 | rknn.config(mean_values=MEAN_VALUES, std_values=STD_VALUES, target_platform=PLATFORM) 48 | 49 | # Load onnx model 50 | print() 51 | print('--- Loading ONNX model ---') 52 | ret = rknn.load_onnx(ONNX_MODEL_PATH) 53 | if ret == 0: 54 | print("Success") 55 | else: 56 | print('Load onnx model failed!') 57 | exit(ret) 58 | 59 | # # Optimize onnx model 60 | # # Only support onnx model that with 'quantization_annotation' information! 61 | # print() 62 | # print('--- Optimize ONNX model ---') 63 | # ret = rknn.optimize_onnx(ONNX_MODEL_PATH, OPT_ONNX_MODEL_PATH) 64 | # if ret == 0: 65 | # print("Success") 66 | # else: 67 | # print('Failed to optimize onnx model!') 68 | 69 | # Build model 70 | print() 71 | print('--- Building model ---') 72 | ret = rknn.build(do_quantization=True, dataset=DATASET) 73 | if ret == 0: 74 | print("Success") 75 | else: 76 | print('Failed to build rknn model.') 77 | exit(ret) 78 | 79 | # Export rknn model 80 | print() 81 | print('--- Export RKNN model: {} ---'.format(RKNN_MODEL_PATH)) 82 | if not os.path.exists(OUT_DIR): 83 | os.makedirs(OUT_DIR) 84 | ret = rknn.export_rknn(RKNN_MODEL_PATH) 85 | if ret == 0: 86 | print("Success") 87 | else: 88 | print('Failed to export rknn model.') 89 | exit(ret) 90 | 91 | # Import Test 92 | print() 93 | print('--- Import Test RKNN model: {} ---'.format(RKNN_MODEL_PATH)) 94 | ret = rknn.load_rknn(RKNN_MODEL_PATH) 95 | if ret == 0: 96 | print("Success") 97 | else: 98 | print('Failed to import rknn model.') 99 | exit(ret) 100 | 101 | print('done') 102 | 103 | rknn.release() 104 | 105 | -------------------------------------------------------------------------------- /rknpu2_ros/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.5) 2 | project(rknpu2_ros) 3 | 4 | # find dependencies 5 | find_package(ament_cmake REQUIRED) 6 | 7 | ament_package() 8 | -------------------------------------------------------------------------------- /rknpu2_ros/README.md: -------------------------------------------------------------------------------- 1 | # RKNPU2_ROS 2 | rknpu2_ros metapackage -------------------------------------------------------------------------------- /rknpu2_ros/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | rknpu2_ros 5 | 0.1.0 6 | rknpu2_ros metapackage 7 | fateshelled 8 | BSD 3-Clause License 9 | fateshelled 10 | 11 | ament_cmake 12 | 13 | rknpu2_ros_common 14 | rknpu2_ros_yolov5 15 | 16 | ament_lint_auto 17 | ament_lint_common 18 | 19 | 20 | ament_cmake 21 | 22 | 23 | -------------------------------------------------------------------------------- /rknpu2_ros_common/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.5) 2 | project(rknpu2_ros_common) 3 | 4 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 5 | add_compile_options(-Wall -Wextra -Wpedantic) 6 | endif() 7 | 8 | # find dependencies 9 | find_package(ament_cmake REQUIRED) 10 | 11 | # library 12 | add_library(${PROJECT_NAME} INTERFACE) 13 | 14 | target_include_directories(${PROJECT_NAME} INTERFACE 15 | $ 16 | $ 17 | ) 18 | 19 | install(TARGETS ${PROJECT_NAME} 20 | EXPORT export_${PROJECT_NAME} 21 | ARCHIVE DESTINATION lib 22 | LIBRARY DESTINATION lib 23 | RUNTIME DESTINATION bin 24 | INCLUDES DESTINATION include 25 | ) 26 | 27 | install( 28 | DIRECTORY include/ 29 | DESTINATION include 30 | ) 31 | 32 | ament_export_targets(export_${PROJECT_NAME}) 33 | 34 | if(BUILD_TESTING) 35 | find_package(ament_lint_auto REQUIRED) 36 | # the following line skips the linter which checks for copyrights 37 | # uncomment the line when a copyright and license is not present in all source files 38 | #set(ament_cmake_copyright_FOUND TRUE) 39 | # the following line skips cpplint (only works in a git repo) 40 | # uncomment the line when this package is not in a git repo 41 | #set(ament_cmake_cpplint_FOUND TRUE) 42 | ament_lint_auto_find_test_dependencies() 43 | endif() 44 | 45 | ament_package() 46 | -------------------------------------------------------------------------------- /rknpu2_ros_common/README.md: -------------------------------------------------------------------------------- 1 | # RKNPU2_ROS_COMMON 2 | 3 | header only package 4 | -------------------------------------------------------------------------------- /rknpu2_ros_common/include/rknpu2_ros_common/coco_names.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | namespace rknpu2_ros{ 7 | static const std::vector COCO_CLASSES = { 8 | "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", 9 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", 10 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 11 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", 12 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 13 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", 14 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", 15 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", 16 | "hair drier", "toothbrush" 17 | }; 18 | const float color_list[80][3] ={ 19 | {0.000, 0.447, 0.741}, 20 | {0.850, 0.325, 0.098}, 21 | {0.929, 0.694, 0.125}, 22 | {0.494, 0.184, 0.556}, 23 | {0.466, 0.674, 0.188}, 24 | {0.301, 0.745, 0.933}, 25 | {0.635, 0.078, 0.184}, 26 | {0.300, 0.300, 0.300}, 27 | {0.600, 0.600, 0.600}, 28 | {1.000, 0.000, 0.000}, 29 | {1.000, 0.500, 0.000}, 30 | {0.749, 0.749, 0.000}, 31 | {0.000, 1.000, 0.000}, 32 | {0.000, 0.000, 1.000}, 33 | {0.667, 0.000, 1.000}, 34 | {0.333, 0.333, 0.000}, 35 | {0.333, 0.667, 0.000}, 36 | {0.333, 1.000, 0.000}, 37 | {0.667, 0.333, 0.000}, 38 | {0.667, 0.667, 0.000}, 39 | {0.667, 1.000, 0.000}, 40 | {1.000, 0.333, 0.000}, 41 | {1.000, 0.667, 0.000}, 42 | {1.000, 1.000, 0.000}, 43 | {0.000, 0.333, 0.500}, 44 | {0.000, 0.667, 0.500}, 45 | {0.000, 1.000, 0.500}, 46 | {0.333, 0.000, 0.500}, 47 | {0.333, 0.333, 0.500}, 48 | {0.333, 0.667, 0.500}, 49 | {0.333, 1.000, 0.500}, 50 | {0.667, 0.000, 0.500}, 51 | {0.667, 0.333, 0.500}, 52 | {0.667, 0.667, 0.500}, 53 | {0.667, 1.000, 0.500}, 54 | {1.000, 0.000, 0.500}, 55 | {1.000, 0.333, 0.500}, 56 | {1.000, 0.667, 0.500}, 57 | {1.000, 1.000, 0.500}, 58 | {0.000, 0.333, 1.000}, 59 | {0.000, 0.667, 1.000}, 60 | {0.000, 1.000, 1.000}, 61 | {0.333, 0.000, 1.000}, 62 | {0.333, 0.333, 1.000}, 63 | {0.333, 0.667, 1.000}, 64 | {0.333, 1.000, 1.000}, 65 | {0.667, 0.000, 1.000}, 66 | {0.667, 0.333, 1.000}, 67 | {0.667, 0.667, 1.000}, 68 | {0.667, 1.000, 1.000}, 69 | {1.000, 0.000, 1.000}, 70 | {1.000, 0.333, 1.000}, 71 | {1.000, 0.667, 1.000}, 72 | {0.333, 0.000, 0.000}, 73 | {0.500, 0.000, 0.000}, 74 | {0.667, 0.000, 0.000}, 75 | {0.833, 0.000, 0.000}, 76 | {1.000, 0.000, 0.000}, 77 | {0.000, 0.167, 0.000}, 78 | {0.000, 0.333, 0.000}, 79 | {0.000, 0.500, 0.000}, 80 | {0.000, 0.667, 0.000}, 81 | {0.000, 0.833, 0.000}, 82 | {0.000, 1.000, 0.000}, 83 | {0.000, 0.000, 0.167}, 84 | {0.000, 0.000, 0.333}, 85 | {0.000, 0.000, 0.500}, 86 | {0.000, 0.000, 0.667}, 87 | {0.000, 0.000, 0.833}, 88 | {0.000, 0.000, 1.000}, 89 | {0.000, 0.000, 0.000}, 90 | {0.143, 0.143, 0.143}, 91 | {0.286, 0.286, 0.286}, 92 | {0.429, 0.429, 0.429}, 93 | {0.571, 0.571, 0.571}, 94 | {0.714, 0.714, 0.714}, 95 | {0.857, 0.857, 0.857}, 96 | {0.000, 0.447, 0.741}, 97 | {0.314, 0.717, 0.741}, 98 | {0.50, 0.5, 0} 99 | }; 100 | } 101 | -------------------------------------------------------------------------------- /rknpu2_ros_common/include/rknpu2_ros_common/drm_func.h: -------------------------------------------------------------------------------- 1 | #ifndef __DRM_FUNC_H__ 2 | #define __DRM_FUNC_H__ 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include // open function 8 | #include // close function 9 | #include 10 | #include 11 | 12 | 13 | #include 14 | #include "libdrm/drm_fourcc.h" 15 | #include "xf86drm.h" 16 | 17 | #ifdef __cplusplus 18 | extern "C" { 19 | #endif 20 | 21 | typedef int (* FUNC_DRM_IOCTL)(int fd, unsigned long request, void *arg); 22 | 23 | typedef struct _drm_context{ 24 | void *drm_handle; 25 | FUNC_DRM_IOCTL io_func; 26 | } drm_context; 27 | 28 | /* memory type definitions. */ 29 | enum drm_rockchip_gem_mem_type 30 | { 31 | /* Physically Continuous memory and used as default. */ 32 | ROCKCHIP_BO_CONTIG = 1 << 0, 33 | /* cachable mapping. */ 34 | ROCKCHIP_BO_CACHABLE = 1 << 1, 35 | /* write-combine mapping. */ 36 | ROCKCHIP_BO_WC = 1 << 2, 37 | ROCKCHIP_BO_SECURE = 1 << 3, 38 | ROCKCHIP_BO_MASK = ROCKCHIP_BO_CONTIG | ROCKCHIP_BO_CACHABLE | 39 | ROCKCHIP_BO_WC | ROCKCHIP_BO_SECURE 40 | }; 41 | 42 | int drm_init(drm_context *drm_ctx); 43 | 44 | void* drm_buf_alloc(drm_context *drm_ctx,int drm_fd, int TexWidth, int TexHeight,int bpp,int *fd,unsigned int *handle,size_t *actual_size); 45 | 46 | int drm_buf_destroy(drm_context *drm_ctx,int drm_fd,int buf_fd, int handle,void *drm_buf,size_t size); 47 | 48 | void drm_deinit(drm_context *drm_ctx, int drm_fd); 49 | 50 | #ifdef __cplusplus 51 | } 52 | #endif 53 | #endif /*__DRM_FUNC_H__*/ -------------------------------------------------------------------------------- /rknpu2_ros_common/include/rknpu2_ros_common/model_loader.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | static unsigned char* load_data(FILE* fp, size_t ofst, size_t sz) 10 | { 11 | unsigned char* data; 12 | int ret; 13 | 14 | data = NULL; 15 | 16 | if (NULL == fp) { 17 | return NULL; 18 | } 19 | 20 | ret = fseek(fp, ofst, SEEK_SET); 21 | if (ret != 0) { 22 | printf("blob seek failure.\n"); 23 | return NULL; 24 | } 25 | 26 | data = (unsigned char*)malloc(sz); 27 | if (data == NULL) { 28 | printf("buffer malloc failure.\n"); 29 | return NULL; 30 | } 31 | ret = fread(data, 1, sz, fp); 32 | return data; 33 | } 34 | 35 | static unsigned char* load_model(const char* filename, int* model_size) 36 | { 37 | FILE* fp; 38 | unsigned char* data; 39 | 40 | fp = fopen(filename, "rb"); 41 | if (NULL == fp) { 42 | printf("Open file %s failed.\n", filename); 43 | return NULL; 44 | } 45 | 46 | fseek(fp, 0, SEEK_END); 47 | int size = ftell(fp); 48 | 49 | data = load_data(fp, 0, size); 50 | 51 | fclose(fp); 52 | 53 | *model_size = size; 54 | return data; 55 | } 56 | 57 | -------------------------------------------------------------------------------- /rknpu2_ros_common/include/rknpu2_ros_common/object.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | namespace rknpu2_ros 6 | { 7 | 8 | struct Object 9 | { 10 | cv::Rect_ rect; 11 | int label; 12 | float prob; 13 | }; 14 | 15 | } -------------------------------------------------------------------------------- /rknpu2_ros_common/include/rknpu2_ros_common/rga_func.h: -------------------------------------------------------------------------------- 1 | #ifndef __RGA_FUNC_H__ 2 | #define __RGA_FUNC_H__ 3 | 4 | #include 5 | #include "RgaApi.h" 6 | 7 | #ifdef __cplusplus 8 | extern "C" { 9 | #endif 10 | 11 | typedef int(* FUNC_RGA_INIT)(); 12 | typedef void(* FUNC_RGA_DEINIT)(); 13 | typedef int(* FUNC_RGA_BLIT)(rga_info_t *, rga_info_t *, rga_info_t *); 14 | 15 | typedef struct _rga_context{ 16 | void *rga_handle; 17 | FUNC_RGA_INIT init_func; 18 | FUNC_RGA_DEINIT deinit_func; 19 | FUNC_RGA_BLIT blit_func; 20 | } rga_context; 21 | 22 | int RGA_init(rga_context* rga_ctx); 23 | 24 | void img_resize_fast(rga_context *rga_ctx, int src_fd, int src_w, int src_h, uint64_t dst_phys, int dst_w, int dst_h); 25 | 26 | void img_resize_slow(rga_context *rga_ctx, void *src_virt, int src_w, int src_h, void *dst_virt, int dst_w, int dst_h); 27 | 28 | int RGA_deinit(rga_context* rga_ctx); 29 | 30 | #ifdef __cplusplus 31 | } 32 | #endif 33 | #endif/*__RGA_FUNC_H__*/ 34 | -------------------------------------------------------------------------------- /rknpu2_ros_common/include/rknpu2_ros_common/utils.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include "object.hpp" 8 | #include "coco_names.hpp" 9 | 10 | #include "bboxes_ex_msgs/msg/bounding_box.hpp" 11 | #include "bboxes_ex_msgs/msg/bounding_boxes.hpp" 12 | 13 | namespace rknpu2_ros 14 | { 15 | namespace utils 16 | { 17 | 18 | static std::vector read_class_labels_file(std::string file_name) 19 | { 20 | std::vector class_names; 21 | std::ifstream ifs(file_name); 22 | std::string buff; 23 | if (ifs.fail()) 24 | { 25 | return class_names; 26 | } 27 | while (getline(ifs, buff)) 28 | { 29 | if (buff == "") 30 | continue; 31 | class_names.push_back(buff); 32 | } 33 | return class_names; 34 | } 35 | 36 | static void draw_objects(cv::Mat bgr, const std::vector &objects, const std::vector &class_names = COCO_CLASSES) 37 | { 38 | 39 | for (size_t i = 0; i < objects.size(); i++) 40 | { 41 | const Object &obj = objects[i]; 42 | 43 | int color_index = obj.label % 80; 44 | cv::Scalar color = cv::Scalar(color_list[color_index][0], color_list[color_index][1], color_list[color_index][2]); 45 | float c_mean = cv::mean(color)[0]; 46 | cv::Scalar txt_color; 47 | if (c_mean > 0.5) 48 | { 49 | txt_color = cv::Scalar(0, 0, 0); 50 | } 51 | else 52 | { 53 | txt_color = cv::Scalar(255, 255, 255); 54 | } 55 | 56 | cv::rectangle(bgr, obj.rect, color * 255, 2); 57 | 58 | char text[256]; 59 | sprintf(text, "%s %.1f%%", class_names[obj.label].c_str(), obj.prob * 100); 60 | 61 | int baseLine = 0; 62 | cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine); 63 | 64 | cv::Scalar txt_bk_color = color * 0.7 * 255; 65 | 66 | int x = obj.rect.x; 67 | int y = obj.rect.y + 1; 68 | if (y > bgr.rows) 69 | y = bgr.rows; 70 | 71 | cv::rectangle(bgr, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)), 72 | txt_bk_color, -1); 73 | 74 | cv::putText(bgr, text, cv::Point(x, y + label_size.height), 75 | cv::FONT_HERSHEY_SIMPLEX, 0.4, txt_color, 1); 76 | } 77 | } 78 | 79 | bboxes_ex_msgs::msg::BoundingBoxes objects_to_bboxes(cv::Mat frame, std_msgs::msg::Header header, std::vector objects, const std::vector labels) 80 | { 81 | bboxes_ex_msgs::msg::BoundingBoxes boxes; 82 | boxes.header = header; 83 | for (auto obj : objects) 84 | { 85 | bboxes_ex_msgs::msg::BoundingBox box; 86 | box.probability = obj.prob; 87 | box.class_id = labels[obj.label]; 88 | box.xmin = obj.rect.x; 89 | box.ymin = obj.rect.y; 90 | box.xmax = (obj.rect.x + obj.rect.width); 91 | box.ymax = (obj.rect.y + obj.rect.height); 92 | box.img_width = frame.cols; 93 | box.img_height = frame.rows; 94 | // tracking id 95 | // box.id = 0; 96 | // depth 97 | // box.center_dist = 0; 98 | boxes.bounding_boxes.emplace_back(box); 99 | } 100 | return boxes; 101 | } 102 | 103 | } 104 | } 105 | -------------------------------------------------------------------------------- /rknpu2_ros_common/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | rknpu2_ros_common 5 | 0.1.0 6 | rknpu2_ros_common package 7 | fateshelled 8 | BSD 3-Clause License 9 | fateshelled 10 | 11 | ament_cmake 12 | 13 | ament_lint_auto 14 | ament_lint_common 15 | 16 | 17 | ament_cmake 18 | 19 | 20 | -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.5) 2 | project(rknpu2_ros_yolov5) 3 | 4 | # Default to C99 5 | if(NOT CMAKE_C_STANDARD) 6 | set(CMAKE_C_STANDARD 99) 7 | endif() 8 | 9 | # Default to C++17 10 | if(NOT CMAKE_CXX_STANDARD) 11 | set(CMAKE_CXX_STANDARD 17) 12 | endif() 13 | 14 | set(CMAKE_CC_COMPILER aarch64-linux-gnu-gcc) 15 | set(CMAKE_CXX_COMPILER aarch64-linux-gnu-g++) 16 | 17 | if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") 18 | add_compile_options(-Wall -Wextra -Wpedantic) 19 | endif() 20 | 21 | set(CMAKE_EXPORT_COMPILE_COMMANDS ON) 22 | 23 | # find dependencies 24 | find_package(ament_cmake REQUIRED) 25 | find_package(rclcpp REQUIRED) 26 | find_package(rclcpp_components REQUIRED) 27 | find_package(std_msgs REQUIRED) 28 | find_package(sensor_msgs REQUIRED) 29 | find_package(cv_bridge REQUIRED) 30 | find_package(image_transport REQUIRED) 31 | find_package(OpenCV REQUIRED) 32 | find_package(bboxes_ex_msgs REQUIRED) 33 | find_package(rknpu2_ros_common REQUIRED) 34 | 35 | # CMAKE_SYSTEM_NAME: Linux 36 | # CMAKE_SYSTEM_PROCESSOR: aarch64 or armhf 37 | # TARGET_SOC: RK3588 or RK356X or RV110X 38 | set(RGA_PATH ${CMAKE_SOURCE_DIR}/../rknpu2/examples/3rdparty/rga/${TARGET_SOC}) 39 | set(RGA_INCLUDE ${RGA_PATH}/include) 40 | set(RGA_LIB ${RGA_PATH}/lib/${CMAKE_SYSTEM_NAME}/${CMAKE_SYSTEM_PROCESSOR}/librga.so) 41 | 42 | set(RKNN_RT_PATH ${CMAKE_SOURCE_DIR}/../rknpu2/runtime/${TARGET_SOC}/${CMAKE_SYSTEM_NAME}/librknn_api) 43 | set(RKNN_RT_LIB ${RKNN_RT_PATH}/${CMAKE_SYSTEM_PROCESSOR}/librknnrt.so) 44 | set(RKNN_RT_INCLUDE ${RKNN_RT_PATH}/include) 45 | 46 | # executable 47 | add_executable(rknpu2_ros_yolov5 48 | src/rknpu2_yolov5.cpp 49 | src/rknpu2_yolov5_node.cpp 50 | ) 51 | 52 | ament_target_dependencies(rknpu2_ros_yolov5 53 | rclcpp 54 | rclcpp_components 55 | cv_bridge 56 | image_transport 57 | builtin_interfaces 58 | std_msgs 59 | sensor_msgs 60 | OpenCV 61 | bboxes_ex_msgs 62 | rknpu2_ros_common 63 | ) 64 | 65 | target_include_directories(rknpu2_ros_yolov5 PUBLIC 66 | $ 67 | $ 68 | ${RGA_INCLUDE} 69 | ${RKNN_RT_INCLUDE} 70 | ) 71 | 72 | target_link_libraries(rknpu2_ros_yolov5 73 | ${RGA_LIB} 74 | ${RKNN_RT_LIB} 75 | ) 76 | 77 | install(TARGETS 78 | rknpu2_ros_yolov5 79 | DESTINATION lib/${PROJECT_NAME} 80 | ) 81 | 82 | install(DIRECTORY 83 | ${CMAKE_SOURCE_DIR}/../rknpu2/examples/rknn_yolov5_demo/model 84 | DESTINATION share/${PROJECT_NAME} 85 | ) 86 | 87 | if(BUILD_TESTING) 88 | find_package(ament_lint_auto REQUIRED) 89 | # the following line skips the linter which checks for copyrights 90 | # uncomment the line when a copyright and license is not present in all source files 91 | #set(ament_cmake_copyright_FOUND TRUE) 92 | # the following line skips cpplint (only works in a git repo) 93 | # uncomment the line when this package is not in a git repo 94 | #set(ament_cmake_cpplint_FOUND TRUE) 95 | ament_lint_auto_find_test_dependencies() 96 | endif() 97 | 98 | ament_package() 99 | -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/README.md: -------------------------------------------------------------------------------- 1 | # rknpu2_ros_yolov5 -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/include/rknpu2_ros_yolov5/postprocess.hpp: -------------------------------------------------------------------------------- 1 | #ifndef RKNPU2_ROS_YOLOV5_POSTPROCESS_HPP_ 2 | #define RKNPU2_ROS_YOLOV5_POSTPROCESS_HPP_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | #include 12 | #include 13 | #include "rknpu2_ros_common/object.hpp" 14 | 15 | namespace rknpu2_ros 16 | { 17 | namespace postprocess 18 | { 19 | const int anchor0[6] = {10, 13, 16, 30, 33, 23}; 20 | const int anchor1[6] = {30, 61, 62, 45, 59, 119}; 21 | const int anchor2[6] = {116, 90, 156, 198, 373, 326}; 22 | 23 | inline static int clamp(float val, int min, int max) { return val > min ? (val < max ? val : max) : min; } 24 | 25 | static float CalculateOverlap(float xmin0, float ymin0, float xmax0, float ymax0, float xmin1, float ymin1, float xmax1, 26 | float ymax1) 27 | { 28 | float w = fmax(0.f, fmin(xmax0, xmax1) - fmax(xmin0, xmin1) + 1.0); 29 | float h = fmax(0.f, fmin(ymax0, ymax1) - fmax(ymin0, ymin1) + 1.0); 30 | float i = w * h; 31 | float u = (xmax0 - xmin0 + 1.0) * (ymax0 - ymin0 + 1.0) + (xmax1 - xmin1 + 1.0) * (ymax1 - ymin1 + 1.0) - i; 32 | return u <= 0.f ? 0.f : (i / u); 33 | } 34 | 35 | static int nms(int validCount, std::vector &outputLocations, std::vector classIds, std::vector &order, 36 | int filterId, float threshold) 37 | { 38 | for (int i = 0; i < validCount; ++i) 39 | { 40 | if (order[i] == -1 || classIds[i] != filterId) 41 | { 42 | continue; 43 | } 44 | int n = order[i]; 45 | for (int j = i + 1; j < validCount; ++j) 46 | { 47 | int m = order[j]; 48 | if (m == -1 || classIds[i] != filterId) 49 | { 50 | continue; 51 | } 52 | float xmin0 = outputLocations[n * 4 + 0]; 53 | float ymin0 = outputLocations[n * 4 + 1]; 54 | float xmax0 = outputLocations[n * 4 + 0] + outputLocations[n * 4 + 2]; 55 | float ymax0 = outputLocations[n * 4 + 1] + outputLocations[n * 4 + 3]; 56 | 57 | float xmin1 = outputLocations[m * 4 + 0]; 58 | float ymin1 = outputLocations[m * 4 + 1]; 59 | float xmax1 = outputLocations[m * 4 + 0] + outputLocations[m * 4 + 2]; 60 | float ymax1 = outputLocations[m * 4 + 1] + outputLocations[m * 4 + 3]; 61 | 62 | float iou = CalculateOverlap(xmin0, ymin0, xmax0, ymax0, xmin1, ymin1, xmax1, ymax1); 63 | 64 | if (iou > threshold) 65 | { 66 | order[j] = -1; 67 | } 68 | } 69 | } 70 | return 0; 71 | } 72 | 73 | static int quick_sort_indice_inverse(std::vector &input, int left, int right, std::vector &indices) 74 | { 75 | float key; 76 | int key_index; 77 | int low = left; 78 | int high = right; 79 | if (left < right) 80 | { 81 | key_index = indices[left]; 82 | key = input[left]; 83 | while (low < high) 84 | { 85 | while (low < high && input[high] <= key) 86 | { 87 | high--; 88 | } 89 | input[low] = input[high]; 90 | indices[low] = indices[high]; 91 | while (low < high && input[low] >= key) 92 | { 93 | low++; 94 | } 95 | input[high] = input[low]; 96 | indices[high] = indices[low]; 97 | } 98 | input[low] = key; 99 | indices[low] = key_index; 100 | quick_sort_indice_inverse(input, left, low - 1, indices); 101 | quick_sort_indice_inverse(input, low + 1, right, indices); 102 | } 103 | return low; 104 | } 105 | 106 | static float sigmoid(float x) { return 1.0 / (1.0 + expf(-x)); } 107 | 108 | static float unsigmoid(float y) { return -1.0 * logf((1.0 / y) - 1.0); } 109 | 110 | inline static int32_t __clip(float val, float min, float max) 111 | { 112 | float f = val <= min ? min : (val >= max ? max : val); 113 | return f; 114 | } 115 | 116 | static int8_t qnt_f32_to_affine(float f32, int32_t zp, float scale) 117 | { 118 | float dst_val = (f32 / scale) + zp; 119 | int8_t res = (int8_t)__clip(dst_val, -128, 127); 120 | return res; 121 | } 122 | 123 | static float deqnt_affine_to_f32(int8_t qnt, int32_t zp, float scale) { return ((float)qnt - (float)zp) * scale; } 124 | 125 | static int process(int8_t *input, int *anchor, int grid_h, int grid_w, int height, int width, int stride, 126 | std::vector &boxes, std::vector &objProbs, std::vector &classId, float threshold, 127 | int32_t zp, float scale, int class_num) 128 | { 129 | int validCount = 0; 130 | int grid_len = grid_h * grid_w; 131 | float thres = unsigmoid(threshold); 132 | int8_t thres_i8 = qnt_f32_to_affine(thres, zp, scale); 133 | for (int a = 0; a < 3; a++) 134 | { 135 | for (int i = 0; i < grid_h; i++) 136 | { 137 | for (int j = 0; j < grid_w; j++) 138 | { 139 | int8_t box_confidence = input[((class_num + 5) * a + 4) * grid_len + i * grid_w + j]; 140 | if (box_confidence >= thres_i8) 141 | { 142 | int offset = ((class_num + 5) * a) * grid_len + i * grid_w + j; 143 | int8_t *in_ptr = input + offset; 144 | float box_x = sigmoid(deqnt_affine_to_f32(*in_ptr, zp, scale)) * 2.0 - 0.5; 145 | float box_y = sigmoid(deqnt_affine_to_f32(in_ptr[grid_len], zp, scale)) * 2.0 - 0.5; 146 | float box_w = sigmoid(deqnt_affine_to_f32(in_ptr[2 * grid_len], zp, scale)) * 2.0; 147 | float box_h = sigmoid(deqnt_affine_to_f32(in_ptr[3 * grid_len], zp, scale)) * 2.0; 148 | box_x = (box_x + j) * (float)stride; 149 | box_y = (box_y + i) * (float)stride; 150 | box_w = box_w * box_w * (float)anchor[a * 2]; 151 | box_h = box_h * box_h * (float)anchor[a * 2 + 1]; 152 | box_x -= (box_w / 2.0); 153 | box_y -= (box_h / 2.0); 154 | 155 | int8_t maxClassProbs = in_ptr[5 * grid_len]; 156 | int maxClassId = 0; 157 | for (int k = 1; k < class_num; ++k) 158 | { 159 | int8_t prob = in_ptr[(5 + k) * grid_len]; 160 | if (prob > maxClassProbs) 161 | { 162 | maxClassId = k; 163 | maxClassProbs = prob; 164 | } 165 | } 166 | if (maxClassProbs > thres_i8) 167 | { 168 | objProbs.push_back(sigmoid(deqnt_affine_to_f32(maxClassProbs, zp, scale)) * sigmoid(deqnt_affine_to_f32(box_confidence, zp, scale))); 169 | classId.push_back(maxClassId); 170 | validCount++; 171 | boxes.push_back(box_x); 172 | boxes.push_back(box_y); 173 | boxes.push_back(box_w); 174 | boxes.push_back(box_h); 175 | } 176 | } 177 | } 178 | } 179 | } 180 | return validCount; 181 | } 182 | 183 | int execute(int8_t *input0, int8_t *input1, int8_t *input2, int model_in_h, int model_in_w, float conf_threshold, 184 | float nms_threshold, float scale_w, float scale_h, std::vector &qnt_zps, 185 | std::vector &qnt_scales, std::vector& results, int num_classes) 186 | { 187 | std::vector filterBoxes; 188 | std::vector objProbs; 189 | std::vector classId; 190 | 191 | // stride 8 192 | int stride0 = 8; 193 | int grid_h0 = model_in_h / stride0; 194 | int grid_w0 = model_in_w / stride0; 195 | int validCount0 = 0; 196 | validCount0 = process(input0, (int *)anchor0, grid_h0, grid_w0, model_in_h, model_in_w, stride0, filterBoxes, objProbs, 197 | classId, conf_threshold, qnt_zps[0], qnt_scales[0], num_classes); 198 | 199 | // stride 16 200 | int stride1 = 16; 201 | int grid_h1 = model_in_h / stride1; 202 | int grid_w1 = model_in_w / stride1; 203 | int validCount1 = 0; 204 | validCount1 = process(input1, (int *)anchor1, grid_h1, grid_w1, model_in_h, model_in_w, stride1, filterBoxes, objProbs, 205 | classId, conf_threshold, qnt_zps[1], qnt_scales[1], num_classes); 206 | 207 | // stride 32 208 | int stride2 = 32; 209 | int grid_h2 = model_in_h / stride2; 210 | int grid_w2 = model_in_w / stride2; 211 | int validCount2 = 0; 212 | validCount2 = process(input2, (int *)anchor2, grid_h2, grid_w2, model_in_h, model_in_w, stride2, filterBoxes, objProbs, 213 | classId, conf_threshold, qnt_zps[2], qnt_scales[2], num_classes); 214 | 215 | int validCount = validCount0 + validCount1 + validCount2; 216 | // no object detect 217 | if (validCount <= 0) 218 | { 219 | return 0; 220 | } 221 | 222 | std::vector indexArray; 223 | for (int i = 0; i < validCount; ++i) 224 | { 225 | indexArray.push_back(i); 226 | } 227 | 228 | quick_sort_indice_inverse(objProbs, 0, validCount - 1, indexArray); 229 | 230 | std::set class_set(std::begin(classId), std::end(classId)); 231 | 232 | for (auto c : class_set) 233 | { 234 | nms(validCount, filterBoxes, classId, indexArray, c, nms_threshold); 235 | } 236 | 237 | int last_count = 0; 238 | /* box valid detect target */ 239 | for (int i = 0; i < validCount; ++i) 240 | { 241 | if (indexArray[i] == -1) 242 | { 243 | continue; 244 | } 245 | int n = indexArray[i]; 246 | 247 | float x1 = filterBoxes[n * 4 + 0]; 248 | float y1 = filterBoxes[n * 4 + 1]; 249 | float x2 = x1 + filterBoxes[n * 4 + 2]; 250 | float y2 = y1 + filterBoxes[n * 4 + 3]; 251 | int id = classId[n]; 252 | float obj_conf = objProbs[i]; 253 | 254 | Object result; 255 | float left = (clamp(x1, 0, model_in_w) / scale_w); 256 | float top = (clamp(y1, 0, model_in_h) / scale_h); 257 | float right = (clamp(x2, 0, model_in_w) / scale_w); 258 | float bottom = (clamp(y2, 0, model_in_h) / scale_h); 259 | result.rect.x = left; 260 | result.rect.y = top; 261 | result.rect.width = right - left; 262 | result.rect.height = bottom - top; 263 | result.prob = obj_conf; 264 | result.label = id; 265 | 266 | results.push_back(result); 267 | 268 | last_count++; 269 | } 270 | 271 | return 0; 272 | } 273 | 274 | } 275 | } 276 | 277 | #endif // RKNPU2_ROS_YOLOV5_POSTPROCESS_HPP_ -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/include/rknpu2_ros_yolov5/rknpu2_yolov5.hpp: -------------------------------------------------------------------------------- 1 | #ifndef RKNPU2_ROS_YOLOV5_HPP_ 2 | #define RKNPU2_ROS_YOLOV5_HPP_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include "RgaUtils.h" 11 | #include "im2d.h" 12 | #include "opencv4/opencv2/opencv.hpp" 13 | #include "rga.h" 14 | #include "rknn_api.h" 15 | 16 | #include "rknpu2_ros_common/object.hpp" 17 | #include "rknpu2_ros_common/model_loader.h" 18 | 19 | namespace rknpu2_ros{ 20 | 21 | class RKNPU2_YoloV5{ 22 | public: 23 | RKNPU2_YoloV5(std::string model_path, 24 | float nms_th=0.45, float conf_th=0.3, 25 | int num_classes=80); 26 | ~RKNPU2_YoloV5(); 27 | 28 | std::vector inference(const cv::Mat& bgr); 29 | private: 30 | std::string model_path_; 31 | float nms_th_; 32 | float conf_th_; 33 | int num_classes_; 34 | 35 | const int anchor0[6] = {10, 13, 16, 30, 33, 23}; 36 | const int anchor1[6] = {30, 61, 62, 45, 59, 119}; 37 | const int anchor2[6] = {116, 90, 156, 198, 373, 326}; 38 | 39 | rknn_context ctx_; 40 | int model_data_size_ = 0; 41 | unsigned char* model_data_; 42 | rga_buffer_t src_; 43 | rga_buffer_t dst_; 44 | im_rect src_rect_; 45 | im_rect dst_rect_; 46 | 47 | rknn_input_output_num io_num_; 48 | rknn_tensor_attr* output_attrs_; 49 | int input_c_; 50 | int input_w_; 51 | int input_h_; 52 | }; 53 | 54 | } // namespace rknpu2_ros 55 | 56 | #endif // RKNPU2_ROS_YOLOV5_HPP_ 57 | -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/include/rknpu2_ros_yolov5/rknpu2_yolov5_node.hpp: -------------------------------------------------------------------------------- 1 | #ifndef RKNPU2_ROS_YOLOV5_NODE_HPP_ 2 | #define RKNPU2_ROS_YOLOV5_NODE_HPP_ 3 | 4 | #include 5 | #include 6 | 7 | #include 8 | 9 | #include 10 | #include 11 | 12 | #include "bboxes_ex_msgs/msg/bounding_box.hpp" 13 | #include "bboxes_ex_msgs/msg/bounding_boxes.hpp" 14 | #include "sensor_msgs/msg/image.hpp" 15 | 16 | #include "rknpu2_ros_yolov5/rknpu2_yolov5.hpp" 17 | 18 | #include "rknpu2_ros_common/coco_names.hpp" 19 | #include "rknpu2_ros_common/utils.hpp" 20 | 21 | namespace rknpu2_ros 22 | { 23 | 24 | class RKNPU2_YoloV5_Node : public rclcpp::Node 25 | { 26 | public: 27 | RKNPU2_YoloV5_Node(const rclcpp::NodeOptions& options); 28 | RKNPU2_YoloV5_Node(const std::string &node_name, const rclcpp::NodeOptions& options); 29 | ~RKNPU2_YoloV5_Node(); 30 | 31 | private: 32 | std::shared_ptr model_; 33 | void initializeParameter(); 34 | void loadLabel(const std::string label_filename); 35 | void image_callback(const sensor_msgs::msg::Image::ConstSharedPtr& ptr); 36 | 37 | image_transport::Subscriber sub_image_; 38 | rclcpp::Publisher::SharedPtr pub_bboxes_; 39 | image_transport::Publisher pub_image_; 40 | 41 | std::vector labels_ = COCO_CLASSES; 42 | 43 | std::string model_path_; 44 | float nms_th_; 45 | float conf_th_; 46 | int num_classes_; 47 | std::string label_path_; 48 | bool imshow_; 49 | std::string src_image_topic_name_; 50 | std::string publish_image_topic_name_; 51 | std::string publish_boundingbox_topic_name_; 52 | 53 | const std::string WINDOW_NAME_ = "RKNPU2_ROS_YOLOV5"; 54 | 55 | }; 56 | 57 | } // namespace rknpu2_ros 58 | 59 | #endif // RKNPU2_ROS_YOLOV5_NODE_HPP_ -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | rknpu2_ros_yolov5 5 | 0.1.0 6 | rknpu2_ros_yolov5 package 7 | fateshelled 8 | Apache-2.0 License 9 | fateshelled 10 | 11 | ament_cmake 12 | 13 | rclcpp 14 | rclcpp_components 15 | OpenCV 16 | cv_bridge 17 | image_transport 18 | std_msgs 19 | sensor_msgs 20 | bboxes_ex_msgs 21 | rknpu2_ros_common 22 | 23 | ament_lint_auto 24 | ament_lint_common 25 | 26 | 27 | ament_cmake 28 | 29 | 30 | -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/src/main.cc: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021 by Rockchip Electronics Co., Ltd. All Rights Reserved. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | /*------------------------------------------- 16 | Includes 17 | -------------------------------------------*/ 18 | #include 19 | #include 20 | #include 21 | #include 22 | #include 23 | 24 | #define _BASETSD_H 25 | 26 | #include "RgaUtils.h" 27 | #include "im2d.h" 28 | #include "opencv4/opencv2/opencv.hpp" 29 | #include "rknpu2_ros_yolov5/postprocess.h" 30 | #include "rga.h" 31 | #include "rknn_api.h" 32 | 33 | /*------------------------------------------- 34 | Functions 35 | -------------------------------------------*/ 36 | 37 | static void dump_tensor_attr(rknn_tensor_attr* attr) 38 | { 39 | printf(" index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, " 40 | "zp=%d, scale=%f\n", 41 | attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3], 42 | attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type), 43 | get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale); 44 | } 45 | 46 | double __get_us(struct timeval t) { return (t.tv_sec * 1000000 + t.tv_usec); } 47 | 48 | static unsigned char* load_data(FILE* fp, size_t ofst, size_t sz) 49 | { 50 | unsigned char* data; 51 | int ret; 52 | 53 | data = NULL; 54 | 55 | if (NULL == fp) { 56 | return NULL; 57 | } 58 | 59 | ret = fseek(fp, ofst, SEEK_SET); 60 | if (ret != 0) { 61 | printf("blob seek failure.\n"); 62 | return NULL; 63 | } 64 | 65 | data = (unsigned char*)malloc(sz); 66 | if (data == NULL) { 67 | printf("buffer malloc failure.\n"); 68 | return NULL; 69 | } 70 | ret = fread(data, 1, sz, fp); 71 | return data; 72 | } 73 | 74 | static unsigned char* load_model(const char* filename, int* model_size) 75 | { 76 | FILE* fp; 77 | unsigned char* data; 78 | 79 | fp = fopen(filename, "rb"); 80 | if (NULL == fp) { 81 | printf("Open file %s failed.\n", filename); 82 | return NULL; 83 | } 84 | 85 | fseek(fp, 0, SEEK_END); 86 | int size = ftell(fp); 87 | 88 | data = load_data(fp, 0, size); 89 | 90 | fclose(fp); 91 | 92 | *model_size = size; 93 | return data; 94 | } 95 | 96 | 97 | /*------------------------------------------- 98 | Main Functions 99 | -------------------------------------------*/ 100 | int main(int argc, char** argv) 101 | { 102 | int status = 0; 103 | char* model_name = NULL; 104 | rknn_context ctx; 105 | size_t actual_size = 0; 106 | int img_width = 0; 107 | int img_height = 0; 108 | int img_channel = 0; 109 | const float nms_threshold = NMS_THRESH; 110 | const float box_conf_threshold = BOX_THRESH; 111 | struct timeval start_time, stop_time; 112 | int ret; 113 | 114 | // init rga context 115 | rga_buffer_t src; 116 | rga_buffer_t dst; 117 | im_rect src_rect; 118 | im_rect dst_rect; 119 | memset(&src_rect, 0, sizeof(src_rect)); 120 | memset(&dst_rect, 0, sizeof(dst_rect)); 121 | memset(&src, 0, sizeof(src)); 122 | memset(&dst, 0, sizeof(dst)); 123 | 124 | 125 | printf("post process config: box_conf_threshold = %.2f, nms_threshold = %.2f\n", box_conf_threshold, nms_threshold); 126 | 127 | model_name = "install/rknpu2_ros_yolov5/share/rknpu2_ros_yolov5/model/RK3588/yolov5s-640-640.rknn"; 128 | /* Create the neural network */ 129 | printf("Loading mode...\n"); 130 | int model_data_size = 0; 131 | unsigned char* model_data = load_model(model_name, &model_data_size); 132 | ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL); 133 | if (ret < 0) { 134 | printf("rknn_init error ret=%d\n", ret); 135 | return -1; 136 | } 137 | 138 | rknn_sdk_version version; 139 | ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version)); 140 | if (ret < 0) { 141 | printf("rknn_init error ret=%d\n", ret); 142 | return -1; 143 | } 144 | printf("sdk version: %s driver version: %s\n", version.api_version, version.drv_version); 145 | 146 | rknn_input_output_num io_num; 147 | ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num)); 148 | if (ret < 0) { 149 | printf("rknn_init error ret=%d\n", ret); 150 | return -1; 151 | } 152 | printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output); 153 | 154 | rknn_tensor_attr input_attrs[io_num.n_input]; 155 | memset(input_attrs, 0, sizeof(input_attrs)); 156 | for (int i = 0; i < io_num.n_input; i++) { 157 | input_attrs[i].index = i; 158 | ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr)); 159 | if (ret < 0) { 160 | printf("rknn_init error ret=%d\n", ret); 161 | return -1; 162 | } 163 | dump_tensor_attr(&(input_attrs[i])); 164 | } 165 | 166 | rknn_tensor_attr output_attrs[io_num.n_output]; 167 | memset(output_attrs, 0, sizeof(output_attrs)); 168 | for (int i = 0; i < io_num.n_output; i++) { 169 | output_attrs[i].index = i; 170 | ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr)); 171 | dump_tensor_attr(&(output_attrs[i])); 172 | } 173 | 174 | int channel = 3; 175 | int width = 0; 176 | int height = 0; 177 | if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) { 178 | printf("model is NCHW input fmt\n"); 179 | channel = input_attrs[0].dims[1]; 180 | height = input_attrs[0].dims[2]; 181 | width = input_attrs[0].dims[3]; 182 | } else { 183 | printf("model is NHWC input fmt\n"); 184 | height = input_attrs[0].dims[1]; 185 | width = input_attrs[0].dims[2]; 186 | channel = input_attrs[0].dims[3]; 187 | } 188 | 189 | printf("model input height=%d, width=%d, channel=%d\n", height, width, channel); 190 | 191 | rknn_input inputs[1]; 192 | memset(inputs, 0, sizeof(inputs)); 193 | inputs[0].index = 0; 194 | inputs[0].type = RKNN_TENSOR_UINT8; 195 | inputs[0].size = width * height * channel; 196 | inputs[0].fmt = RKNN_TENSOR_NHWC; 197 | inputs[0].pass_through = 0; 198 | 199 | 200 | cv::VideoCapture cap("/dev/video0"); 201 | 202 | cv::Mat img, orig_img; 203 | cap >> orig_img; 204 | cv::cvtColor(orig_img, img, cv::COLOR_BGR2RGB); 205 | img_width = img.cols; 206 | img_height = img.rows; 207 | printf("img width = %d, img height = %d\n", img_width, img_height); 208 | 209 | while(true) 210 | { 211 | cap >> orig_img; 212 | cv::cvtColor(orig_img, img, cv::COLOR_BGR2RGB); 213 | 214 | cv::resize(img, img, cv::Size(width, height)); 215 | inputs[0].buf = (void*)img.data; 216 | 217 | gettimeofday(&start_time, NULL); 218 | rknn_inputs_set(ctx, io_num.n_input, inputs); 219 | 220 | rknn_output outputs[io_num.n_output]; 221 | memset(outputs, 0, sizeof(outputs)); 222 | for (int i = 0; i < io_num.n_output; i++) { 223 | outputs[i].want_float = 0; 224 | } 225 | 226 | ret = rknn_run(ctx, NULL); 227 | ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL); 228 | gettimeofday(&stop_time, NULL); 229 | printf("once run use %f ms\n", (__get_us(stop_time) - __get_us(start_time)) / 1000); 230 | 231 | // post process 232 | float scale_w = (float)width / img_width; 233 | float scale_h = (float)height / img_height; 234 | 235 | detect_result_group_t detect_result_group; 236 | std::vector out_scales; 237 | std::vector out_zps; 238 | for (int i = 0; i < io_num.n_output; ++i) { 239 | out_scales.push_back(output_attrs[i].scale); 240 | out_zps.push_back(output_attrs[i].zp); 241 | } 242 | post_process((int8_t*)outputs[0].buf, (int8_t*)outputs[1].buf, (int8_t*)outputs[2].buf, height, width, 243 | box_conf_threshold, nms_threshold, scale_w, scale_h, out_zps, out_scales, &detect_result_group); 244 | 245 | // Draw Objects 246 | char text[256]; 247 | for (int i = 0; i < detect_result_group.count; i++) { 248 | detect_result_t* det_result = &(detect_result_group.results[i]); 249 | sprintf(text, "%s %.1f%%", det_result->name, det_result->prop * 100); 250 | printf("%s @ (%d %d %d %d) %f\n", det_result->name, det_result->box.left, det_result->box.top, 251 | det_result->box.right, det_result->box.bottom, det_result->prop); 252 | int x1 = det_result->box.left; 253 | int y1 = det_result->box.top; 254 | int x2 = det_result->box.right; 255 | int y2 = det_result->box.bottom; 256 | rectangle(orig_img, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(255, 0, 0, 255), 3); 257 | putText(orig_img, text, cv::Point(x1, y1 + 12), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0)); 258 | } 259 | ret = rknn_outputs_release(ctx, io_num.n_output, outputs); 260 | cv::imshow("yolov5", orig_img); 261 | int key = cv::waitKey(10); 262 | if(key == 27) 263 | { 264 | break; 265 | } 266 | } 267 | cv::destroyAllWindows(); 268 | deinitPostProcess(); 269 | // release 270 | ret = rknn_destroy(ctx); 271 | if (model_data) { 272 | free(model_data); 273 | } 274 | // if (resize_buf) { 275 | // free(resize_buf); 276 | // } 277 | return 0; 278 | } 279 | -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/src/postprocess.cc: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2021 by Rockchip Electronics Co., Ltd. All Rights Reserved. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #include "rknpu2_ros_yolov5/postprocess.h" 16 | 17 | #include 18 | #include 19 | #include 20 | #include 21 | #include 22 | #include 23 | 24 | #include 25 | #include 26 | #define LABEL_NALE_TXT_PATH "install/rknpu2_ros_yolov5/share/rknpu2_ros_yolov5/model/coco_80_labels_list.txt" 27 | 28 | static char* labels[OBJ_CLASS_NUM]; 29 | 30 | const int anchor0[6] = {10, 13, 16, 30, 33, 23}; 31 | const int anchor1[6] = {30, 61, 62, 45, 59, 119}; 32 | const int anchor2[6] = {116, 90, 156, 198, 373, 326}; 33 | 34 | inline static int clamp(float val, int min, int max) { return val > min ? (val < max ? val : max) : min; } 35 | 36 | char* readLine(FILE* fp, char* buffer, int* len) 37 | { 38 | int ch; 39 | int i = 0; 40 | size_t buff_len = 0; 41 | 42 | buffer = (char*)malloc(buff_len + 1); 43 | if (!buffer) 44 | return NULL; // Out of memory 45 | 46 | while ((ch = fgetc(fp)) != '\n' && ch != EOF) { 47 | buff_len++; 48 | void* tmp = realloc(buffer, buff_len + 1); 49 | if (tmp == NULL) { 50 | free(buffer); 51 | return NULL; // Out of memory 52 | } 53 | buffer = (char*)tmp; 54 | 55 | buffer[i] = (char)ch; 56 | i++; 57 | } 58 | buffer[i] = '\0'; 59 | 60 | *len = buff_len; 61 | 62 | // Detect end 63 | if (ch == EOF && (i == 0 || ferror(fp))) { 64 | free(buffer); 65 | return NULL; 66 | } 67 | return buffer; 68 | } 69 | 70 | int readLines(const char* fileName, char* lines[], int max_line) 71 | { 72 | FILE* file = fopen(fileName, "r"); 73 | char* s; 74 | int i = 0; 75 | int n = 0; 76 | 77 | if (file == NULL) { 78 | printf("Open %s fail!\n", fileName); 79 | return -1; 80 | } 81 | 82 | while ((s = readLine(file, s, &n)) != NULL) { 83 | lines[i++] = s; 84 | if (i >= max_line) 85 | break; 86 | } 87 | fclose(file); 88 | return i; 89 | } 90 | 91 | int loadLabelName(const char* locationFilename, char* label[]) 92 | { 93 | printf("loadLabelName %s\n", locationFilename); 94 | readLines(locationFilename, label, OBJ_CLASS_NUM); 95 | return 0; 96 | } 97 | 98 | static float CalculateOverlap(float xmin0, float ymin0, float xmax0, float ymax0, float xmin1, float ymin1, float xmax1, 99 | float ymax1) 100 | { 101 | float w = fmax(0.f, fmin(xmax0, xmax1) - fmax(xmin0, xmin1) + 1.0); 102 | float h = fmax(0.f, fmin(ymax0, ymax1) - fmax(ymin0, ymin1) + 1.0); 103 | float i = w * h; 104 | float u = (xmax0 - xmin0 + 1.0) * (ymax0 - ymin0 + 1.0) + (xmax1 - xmin1 + 1.0) * (ymax1 - ymin1 + 1.0) - i; 105 | return u <= 0.f ? 0.f : (i / u); 106 | } 107 | 108 | static int nms(int validCount, std::vector& outputLocations, std::vector classIds, std::vector& order, 109 | int filterId, float threshold) 110 | { 111 | for (int i = 0; i < validCount; ++i) { 112 | if (order[i] == -1 || classIds[i] != filterId) { 113 | continue; 114 | } 115 | int n = order[i]; 116 | for (int j = i + 1; j < validCount; ++j) { 117 | int m = order[j]; 118 | if (m == -1 || classIds[i] != filterId) { 119 | continue; 120 | } 121 | float xmin0 = outputLocations[n * 4 + 0]; 122 | float ymin0 = outputLocations[n * 4 + 1]; 123 | float xmax0 = outputLocations[n * 4 + 0] + outputLocations[n * 4 + 2]; 124 | float ymax0 = outputLocations[n * 4 + 1] + outputLocations[n * 4 + 3]; 125 | 126 | float xmin1 = outputLocations[m * 4 + 0]; 127 | float ymin1 = outputLocations[m * 4 + 1]; 128 | float xmax1 = outputLocations[m * 4 + 0] + outputLocations[m * 4 + 2]; 129 | float ymax1 = outputLocations[m * 4 + 1] + outputLocations[m * 4 + 3]; 130 | 131 | float iou = CalculateOverlap(xmin0, ymin0, xmax0, ymax0, xmin1, ymin1, xmax1, ymax1); 132 | 133 | if (iou > threshold) { 134 | order[j] = -1; 135 | } 136 | } 137 | } 138 | return 0; 139 | } 140 | 141 | static int quick_sort_indice_inverse(std::vector& input, int left, int right, std::vector& indices) 142 | { 143 | float key; 144 | int key_index; 145 | int low = left; 146 | int high = right; 147 | if (left < right) { 148 | key_index = indices[left]; 149 | key = input[left]; 150 | while (low < high) { 151 | while (low < high && input[high] <= key) { 152 | high--; 153 | } 154 | input[low] = input[high]; 155 | indices[low] = indices[high]; 156 | while (low < high && input[low] >= key) { 157 | low++; 158 | } 159 | input[high] = input[low]; 160 | indices[high] = indices[low]; 161 | } 162 | input[low] = key; 163 | indices[low] = key_index; 164 | quick_sort_indice_inverse(input, left, low - 1, indices); 165 | quick_sort_indice_inverse(input, low + 1, right, indices); 166 | } 167 | return low; 168 | } 169 | 170 | static float sigmoid(float x) { return 1.0 / (1.0 + expf(-x)); } 171 | 172 | static float unsigmoid(float y) { return -1.0 * logf((1.0 / y) - 1.0); } 173 | 174 | inline static int32_t __clip(float val, float min, float max) 175 | { 176 | float f = val <= min ? min : (val >= max ? max : val); 177 | return f; 178 | } 179 | 180 | static int8_t qnt_f32_to_affine(float f32, int32_t zp, float scale) 181 | { 182 | float dst_val = (f32 / scale) + zp; 183 | int8_t res = (int8_t)__clip(dst_val, -128, 127); 184 | return res; 185 | } 186 | 187 | static float deqnt_affine_to_f32(int8_t qnt, int32_t zp, float scale) { return ((float)qnt - (float)zp) * scale; } 188 | 189 | static int process(int8_t* input, int* anchor, int grid_h, int grid_w, int height, int width, int stride, 190 | std::vector& boxes, std::vector& objProbs, std::vector& classId, float threshold, 191 | int32_t zp, float scale) 192 | { 193 | int validCount = 0; 194 | int grid_len = grid_h * grid_w; 195 | float thres = unsigmoid(threshold); 196 | int8_t thres_i8 = qnt_f32_to_affine(thres, zp, scale); 197 | for (int a = 0; a < 3; a++) { 198 | for (int i = 0; i < grid_h; i++) { 199 | for (int j = 0; j < grid_w; j++) { 200 | int8_t box_confidence = input[(PROP_BOX_SIZE * a + 4) * grid_len + i * grid_w + j]; 201 | if (box_confidence >= thres_i8) { 202 | int offset = (PROP_BOX_SIZE * a) * grid_len + i * grid_w + j; 203 | int8_t* in_ptr = input + offset; 204 | float box_x = sigmoid(deqnt_affine_to_f32(*in_ptr, zp, scale)) * 2.0 - 0.5; 205 | float box_y = sigmoid(deqnt_affine_to_f32(in_ptr[grid_len], zp, scale)) * 2.0 - 0.5; 206 | float box_w = sigmoid(deqnt_affine_to_f32(in_ptr[2 * grid_len], zp, scale)) * 2.0; 207 | float box_h = sigmoid(deqnt_affine_to_f32(in_ptr[3 * grid_len], zp, scale)) * 2.0; 208 | box_x = (box_x + j) * (float)stride; 209 | box_y = (box_y + i) * (float)stride; 210 | box_w = box_w * box_w * (float)anchor[a * 2]; 211 | box_h = box_h * box_h * (float)anchor[a * 2 + 1]; 212 | box_x -= (box_w / 2.0); 213 | box_y -= (box_h / 2.0); 214 | 215 | int8_t maxClassProbs = in_ptr[5 * grid_len]; 216 | int maxClassId = 0; 217 | for (int k = 1; k < OBJ_CLASS_NUM; ++k) { 218 | int8_t prob = in_ptr[(5 + k) * grid_len]; 219 | if (prob > maxClassProbs) { 220 | maxClassId = k; 221 | maxClassProbs = prob; 222 | } 223 | } 224 | if (maxClassProbs>thres_i8){ 225 | objProbs.push_back(sigmoid(deqnt_affine_to_f32(maxClassProbs, zp, scale))* sigmoid(deqnt_affine_to_f32(box_confidence, zp, scale))); 226 | classId.push_back(maxClassId); 227 | validCount++; 228 | boxes.push_back(box_x); 229 | boxes.push_back(box_y); 230 | boxes.push_back(box_w); 231 | boxes.push_back(box_h); 232 | } 233 | } 234 | } 235 | } 236 | } 237 | return validCount; 238 | } 239 | 240 | int post_process(int8_t* input0, int8_t* input1, int8_t* input2, int model_in_h, int model_in_w, float conf_threshold, 241 | float nms_threshold, float scale_w, float scale_h, std::vector& qnt_zps, 242 | std::vector& qnt_scales, detect_result_group_t* group) 243 | { 244 | static int init = -1; 245 | if (init == -1) { 246 | int ret = 0; 247 | ret = loadLabelName(LABEL_NALE_TXT_PATH, labels); 248 | if (ret < 0) { 249 | return -1; 250 | } 251 | 252 | init = 0; 253 | } 254 | memset(group, 0, sizeof(detect_result_group_t)); 255 | 256 | std::vector filterBoxes; 257 | std::vector objProbs; 258 | std::vector classId; 259 | 260 | // stride 8 261 | int stride0 = 8; 262 | int grid_h0 = model_in_h / stride0; 263 | int grid_w0 = model_in_w / stride0; 264 | int validCount0 = 0; 265 | validCount0 = process(input0, (int*)anchor0, grid_h0, grid_w0, model_in_h, model_in_w, stride0, filterBoxes, objProbs, 266 | classId, conf_threshold, qnt_zps[0], qnt_scales[0]); 267 | 268 | // stride 16 269 | int stride1 = 16; 270 | int grid_h1 = model_in_h / stride1; 271 | int grid_w1 = model_in_w / stride1; 272 | int validCount1 = 0; 273 | validCount1 = process(input1, (int*)anchor1, grid_h1, grid_w1, model_in_h, model_in_w, stride1, filterBoxes, objProbs, 274 | classId, conf_threshold, qnt_zps[1], qnt_scales[1]); 275 | 276 | // stride 32 277 | int stride2 = 32; 278 | int grid_h2 = model_in_h / stride2; 279 | int grid_w2 = model_in_w / stride2; 280 | int validCount2 = 0; 281 | validCount2 = process(input2, (int*)anchor2, grid_h2, grid_w2, model_in_h, model_in_w, stride2, filterBoxes, objProbs, 282 | classId, conf_threshold, qnt_zps[2], qnt_scales[2]); 283 | 284 | int validCount = validCount0 + validCount1 + validCount2; 285 | // no object detect 286 | if (validCount <= 0) { 287 | return 0; 288 | } 289 | 290 | std::vector indexArray; 291 | for (int i = 0; i < validCount; ++i) { 292 | indexArray.push_back(i); 293 | } 294 | 295 | quick_sort_indice_inverse(objProbs, 0, validCount - 1, indexArray); 296 | 297 | std::set class_set(std::begin(classId), std::end(classId)); 298 | 299 | for (auto c : class_set) { 300 | nms(validCount, filterBoxes, classId, indexArray, c, nms_threshold); 301 | } 302 | 303 | int last_count = 0; 304 | group->count = 0; 305 | /* box valid detect target */ 306 | for (int i = 0; i < validCount; ++i) { 307 | if (indexArray[i] == -1 || last_count >= OBJ_NUMB_MAX_SIZE) { 308 | continue; 309 | } 310 | int n = indexArray[i]; 311 | 312 | float x1 = filterBoxes[n * 4 + 0]; 313 | float y1 = filterBoxes[n * 4 + 1]; 314 | float x2 = x1 + filterBoxes[n * 4 + 2]; 315 | float y2 = y1 + filterBoxes[n * 4 + 3]; 316 | int id = classId[n]; 317 | float obj_conf = objProbs[i]; 318 | 319 | group->results[last_count].box.left = (int)(clamp(x1, 0, model_in_w) / scale_w); 320 | group->results[last_count].box.top = (int)(clamp(y1, 0, model_in_h) / scale_h); 321 | group->results[last_count].box.right = (int)(clamp(x2, 0, model_in_w) / scale_w); 322 | group->results[last_count].box.bottom = (int)(clamp(y2, 0, model_in_h) / scale_h); 323 | group->results[last_count].prop = obj_conf; 324 | char* label = labels[id]; 325 | strncpy(group->results[last_count].name, label, OBJ_NAME_MAX_SIZE); 326 | 327 | // printf("result %2d: (%4d, %4d, %4d, %4d), %s\n", i, group->results[last_count].box.left, 328 | // group->results[last_count].box.top, 329 | // group->results[last_count].box.right, group->results[last_count].box.bottom, label); 330 | last_count++; 331 | } 332 | group->count = last_count; 333 | 334 | return 0; 335 | } 336 | 337 | void deinitPostProcess() 338 | { 339 | for (int i = 0; i < OBJ_CLASS_NUM; i++) { 340 | if (labels[i] != nullptr) { 341 | free(labels[i]); 342 | labels[i] = nullptr; 343 | } 344 | } 345 | } 346 | -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/src/postprocess.h: -------------------------------------------------------------------------------- 1 | #ifndef _RKNN_ZERO_COPY_DEMO_POSTPROCESS_H_ 2 | #define _RKNN_ZERO_COPY_DEMO_POSTPROCESS_H_ 3 | 4 | #include 5 | #include 6 | 7 | #define OBJ_NAME_MAX_SIZE 16 8 | #define OBJ_NUMB_MAX_SIZE 64 9 | #define OBJ_CLASS_NUM 80 10 | #define NMS_THRESH 0.45 11 | #define BOX_THRESH 0.25 12 | #define PROP_BOX_SIZE (5+OBJ_CLASS_NUM) 13 | 14 | typedef struct _BOX_RECT 15 | { 16 | int left; 17 | int right; 18 | int top; 19 | int bottom; 20 | } BOX_RECT; 21 | 22 | typedef struct __detect_result_t 23 | { 24 | char name[OBJ_NAME_MAX_SIZE]; 25 | BOX_RECT box; 26 | float prop; 27 | } detect_result_t; 28 | 29 | typedef struct _detect_result_group_t 30 | { 31 | int id; 32 | int count; 33 | detect_result_t results[OBJ_NUMB_MAX_SIZE]; 34 | } detect_result_group_t; 35 | 36 | int post_process(int8_t *input0, int8_t *input1, int8_t *input2, int model_in_h, int model_in_w, 37 | float conf_threshold, float nms_threshold, float scale_w, float scale_h, 38 | std::vector &qnt_zps, std::vector &qnt_scales, 39 | detect_result_group_t *group); 40 | 41 | void deinitPostProcess(); 42 | #endif //_RKNN_ZERO_COPY_DEMO_POSTPROCESS_H_ 43 | -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/src/rknpu2_yolov5.cpp: -------------------------------------------------------------------------------- 1 | #include "rknpu2_ros_yolov5/rknpu2_yolov5.hpp" 2 | #include "rknpu2_ros_yolov5/postprocess.hpp" 3 | 4 | 5 | namespace rknpu2_ros 6 | { 7 | RKNPU2_YoloV5::RKNPU2_YoloV5(std::string model_path, float nms_th, float conf_th, int num_classes) 8 | : model_path_(model_path), nms_th_(nms_th), conf_th_(conf_th), num_classes_(num_classes) 9 | { 10 | int ret = 0; 11 | memset(&this->src_rect_, 0, sizeof(this->src_rect_)); 12 | memset(&this->dst_rect_, 0, sizeof(this->dst_rect_)); 13 | memset(&this->src_, 0, sizeof(this->src_)); 14 | memset(&this->dst_, 0, sizeof(this->dst_)); 15 | 16 | this->model_data_ = load_model(model_path_.c_str(), &this->model_data_size_); 17 | ret = rknn_init(&this->ctx_, this->model_data_, this->model_data_size_, 0, NULL); 18 | if (ret < 0) { 19 | printf("rknn_init error ret=%d\n", ret); 20 | exit -1; 21 | } 22 | 23 | rknn_sdk_version version; 24 | ret = rknn_query(this->ctx_, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version)); 25 | if (ret < 0) { 26 | printf("rknn_init error ret=%d\n", ret); 27 | exit -1; 28 | } 29 | printf("sdk version: %s driver version: %s\n", version.api_version, version.drv_version); 30 | 31 | // input/output num 32 | ret = rknn_query(this->ctx_, RKNN_QUERY_IN_OUT_NUM, &io_num_, sizeof(io_num_)); 33 | if (ret < 0) { 34 | printf("rknn_init error ret=%d\n", ret); 35 | exit -1; 36 | } 37 | printf("model input num: %d, output num: %d\n", io_num_.n_input, io_num_.n_output); 38 | 39 | rknn_tensor_attr input_attrs[io_num_.n_input]; 40 | memset(input_attrs, 0, sizeof(input_attrs)); 41 | for (int i = 0; i < io_num_.n_input; i++) { 42 | input_attrs[i].index = i; 43 | ret = rknn_query(this->ctx_, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr)); 44 | if (ret < 0) { 45 | printf("rknn_init error ret=%d\n", ret); 46 | exit -1; 47 | } 48 | } 49 | 50 | if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) { 51 | printf("model is NCHW input fmt\n"); 52 | input_c_ = input_attrs[0].dims[1]; 53 | input_h_ = input_attrs[0].dims[2]; 54 | input_w_ = input_attrs[0].dims[3]; 55 | } else { 56 | printf("model is NHWC input fmt\n"); 57 | input_h_ = input_attrs[0].dims[1]; 58 | input_w_ = input_attrs[0].dims[2]; 59 | input_c_ = input_attrs[0].dims[3]; 60 | } 61 | 62 | printf("model input height=%d, width=%d, channel=%d\n", input_h_, input_w_, input_c_); 63 | 64 | } 65 | 66 | RKNPU2_YoloV5::~RKNPU2_YoloV5() 67 | { 68 | rknn_destroy(this->ctx_); 69 | free(this->model_data_); 70 | } 71 | 72 | std::vector RKNPU2_YoloV5::inference(const cv::Mat &bgr) 73 | { 74 | cv::Mat img; 75 | int ret = 0; 76 | 77 | int img_width = bgr.cols; 78 | int img_height = bgr.rows; 79 | 80 | // preprocess 81 | cv::cvtColor(bgr, img, cv::COLOR_BGR2RGB); 82 | cv::resize(img, img, cv::Size(input_w_, input_h_)); 83 | 84 | // set input 85 | rknn_input inputs[1]; 86 | memset(inputs, 0, sizeof(inputs)); 87 | inputs[0].index = 0; 88 | inputs[0].type = RKNN_TENSOR_UINT8; 89 | inputs[0].size = input_w_ * input_h_ * input_c_; 90 | inputs[0].fmt = RKNN_TENSOR_NHWC; 91 | inputs[0].pass_through = 0; 92 | inputs[0].buf = (void*)img.data; 93 | 94 | rknn_inputs_set(this->ctx_, this->io_num_.n_input, inputs); 95 | 96 | rknn_output outputs[this->io_num_.n_output]; 97 | memset(outputs, 0, sizeof(outputs)); 98 | for (int i = 0; i < this->io_num_.n_output; i++) { 99 | outputs[i].want_float = 0; 100 | } 101 | 102 | rknn_tensor_attr output_attrs[this->io_num_.n_output]; 103 | memset(output_attrs, 0, sizeof(output_attrs)); 104 | for (int i = 0; i < io_num_.n_output; i++) { 105 | output_attrs[i].index = i; 106 | ret = rknn_query(this->ctx_, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr)); 107 | } 108 | 109 | // inference 110 | ret = rknn_run(this->ctx_, NULL); 111 | 112 | // get output 113 | ret = rknn_outputs_get(this->ctx_, this->io_num_.n_output, outputs, NULL); 114 | 115 | // post process 116 | float scale_w = (float)input_w_ / img_width; 117 | float scale_h = (float)input_h_ / img_height; 118 | 119 | std::vector results; 120 | std::vector out_scales; 121 | std::vector out_zps; 122 | for (int i = 0; i < this->io_num_.n_output; ++i) { 123 | out_scales.push_back(output_attrs[i].scale); 124 | out_zps.push_back(output_attrs[i].zp); 125 | } 126 | postprocess::execute((int8_t*)outputs[0].buf, (int8_t*)outputs[1].buf, (int8_t*)outputs[2].buf, input_h_, input_w_, 127 | this->conf_th_, this->nms_th_, scale_w, scale_h, out_zps, out_scales, results, this->num_classes_); 128 | rknn_outputs_release(this->ctx_, this->io_num_.n_output, outputs); 129 | return results; 130 | } 131 | 132 | } 133 | -------------------------------------------------------------------------------- /rknpu2_ros_yolov5/src/rknpu2_yolov5_node.cpp: -------------------------------------------------------------------------------- 1 | #include "rknpu2_ros_yolov5/rknpu2_yolov5_node.hpp" 2 | 3 | namespace rknpu2_ros 4 | { 5 | RKNPU2_YoloV5_Node::RKNPU2_YoloV5_Node(const rclcpp::NodeOptions &options) 6 | : RKNPU2_YoloV5_Node::RKNPU2_YoloV5_Node("", options) 7 | {} 8 | 9 | RKNPU2_YoloV5_Node::RKNPU2_YoloV5_Node(const std::string &node_name, const rclcpp::NodeOptions &options) 10 | : rclcpp::Node("rknpu2_yolov5_node", node_name, options) 11 | { 12 | this->initializeParameter(); 13 | 14 | this->model_ = std::make_unique(this->model_path_, this->nms_th_, this->conf_th_, this->num_classes_); 15 | RCLCPP_INFO(this->get_logger(), "model loaded"); 16 | 17 | this->sub_image_ = image_transport::create_subscription( 18 | this, this->src_image_topic_name_, 19 | std::bind(&RKNPU2_YoloV5_Node::image_callback, this, std::placeholders::_1), 20 | "raw"); 21 | this->pub_bboxes_ = this->create_publisher( 22 | this->publish_boundingbox_topic_name_, 23 | 10 24 | ); 25 | this->pub_image_ = image_transport::create_publisher(this, this->publish_image_topic_name_); 26 | } 27 | 28 | RKNPU2_YoloV5_Node::~RKNPU2_YoloV5_Node() 29 | {} 30 | 31 | void RKNPU2_YoloV5_Node::initializeParameter() 32 | { 33 | this->model_path_ = this->declare_parameter("model_path", "install/rknpu2_ros_yolov5/share/rknpu2_ros_yolov5/model/RK3588/yolov5s-640-640.rknn"); 34 | this->label_path_ = this->declare_parameter("label_path", ""); 35 | this->nms_th_ = this->declare_parameter("nms_th", 0.45); 36 | this->conf_th_ = this->declare_parameter("conf_th", 0.30); 37 | this->num_classes_ = this->declare_parameter("num_classes", 80); 38 | this->imshow_ = this->declare_parameter("imshow", true); 39 | this->src_image_topic_name_ = this->declare_parameter("src_image_topic_name", "image_raw"); 40 | this->publish_image_topic_name_ = this->declare_parameter("publish_image_topic_name", "yolov5/image_raw"); 41 | this->publish_boundingbox_topic_name_ = this->declare_parameter("publish_boundingbox_topic_name", "yolov5/bounding_boxes"); 42 | 43 | RCLCPP_INFO(this->get_logger(), "Set parameter model_path: '%s'", this->model_path_.c_str()); 44 | RCLCPP_INFO(this->get_logger(), "Set parameter label_path: '%s'", this->label_path_.c_str()); 45 | RCLCPP_INFO(this->get_logger(), "Set parameter conf_th: %f", this->conf_th_); 46 | RCLCPP_INFO(this->get_logger(), "Set parameter nms_th: %f", this->nms_th_); 47 | RCLCPP_INFO(this->get_logger(), "Set parameter num_classes: %i", this->num_classes_); 48 | RCLCPP_INFO(this->get_logger(), "Set parameter imshow: %i", this->imshow_); 49 | RCLCPP_INFO(this->get_logger(), "Set parameter src_image_topic_name: '%s'", this->src_image_topic_name_.c_str()); 50 | RCLCPP_INFO(this->get_logger(), "Set parameter publish_image_topic_name: '%s'", this->publish_image_topic_name_.c_str()); 51 | RCLCPP_INFO(this->get_logger(), "Set parameter publish_boundingbox_topic_name: '%s'", this->publish_boundingbox_topic_name_.c_str()); 52 | } 53 | 54 | void RKNPU2_YoloV5_Node::image_callback(const sensor_msgs::msg::Image::ConstSharedPtr& ptr) 55 | { 56 | // auto img = cv_bridge::toCvShare(ptr, "bgr8"); 57 | auto img = cv_bridge::toCvCopy(ptr, "bgr8"); 58 | cv::Mat frame = img->image; 59 | 60 | // fps 61 | auto now = std::chrono::system_clock::now(); 62 | 63 | auto objects = this->model_->inference(frame); 64 | 65 | auto end = std::chrono::system_clock::now(); 66 | auto elapsed = std::chrono::duration_cast(end - now); 67 | RCLCPP_INFO(this->get_logger(), "Inference: %f FPS", 1000.0f / elapsed.count()); 68 | 69 | // draw 70 | cv::Mat drawn; 71 | frame.copyTo(drawn); 72 | utils::draw_objects(drawn, objects, this->labels_); 73 | 74 | // imshow 75 | if(this->imshow_){ 76 | cv::imshow(this->WINDOW_NAME_, drawn); 77 | auto key = cv::waitKey(1); 78 | if(key == 27){ 79 | rclcpp::shutdown(); 80 | } 81 | } 82 | 83 | // pub bbox 84 | auto bboxes = utils::objects_to_bboxes(frame, ptr->header, objects, this->labels_); 85 | this->pub_bboxes_->publish(bboxes); 86 | 87 | // pub img 88 | sensor_msgs::msg::Image::SharedPtr pub_img; 89 | pub_img = cv_bridge::CvImage(img->header, "bgr8", drawn).toImageMsg(); 90 | this->pub_image_.publish(pub_img); 91 | } 92 | 93 | } 94 | 95 | int main(int argc, char *argv[]) 96 | { 97 | rclcpp::init(argc, argv); 98 | rclcpp::NodeOptions node_options; 99 | rclcpp::spin(std::make_shared(node_options)); 100 | rclcpp::shutdown(); 101 | return 0; 102 | } 103 | 104 | // #include 105 | // RCLCPP_COMPONENTS_REGISTER_NODE(rknpu2_ros::RKNPU2_YoloV5_Node) --------------------------------------------------------------------------------