开始第三阶段开发
This commit is contained in:
parent
4cd65bd168
commit
ad1de5461c
76
configs/sample_ai_pipeline.json
Normal file
76
configs/sample_ai_pipeline.json
Normal file
@ -0,0 +1,76 @@
|
||||
{
|
||||
"queue": { "size": 8, "strategy": "drop_oldest" },
|
||||
"graphs": [
|
||||
{
|
||||
"name": "cam1_ai_pipeline",
|
||||
"nodes": [
|
||||
{
|
||||
"id": "in_cam1",
|
||||
"type": "input_rtsp",
|
||||
"role": "source",
|
||||
"enable": true,
|
||||
"url": "rtsp://10.0.0.9:8554/cam",
|
||||
"fps": 25,
|
||||
"width": 1920,
|
||||
"height": 1080,
|
||||
"use_mpp": false,
|
||||
"use_ffmpeg": true,
|
||||
"force_tcp": true
|
||||
},
|
||||
{
|
||||
"id": "pre_cam1",
|
||||
"type": "preprocess",
|
||||
"role": "filter",
|
||||
"enable": true,
|
||||
"dst_w": 640,
|
||||
"dst_h": 640,
|
||||
"dst_format": "rgb",
|
||||
"keep_ratio": false,
|
||||
"use_rga": true
|
||||
},
|
||||
{
|
||||
"id": "ai_cam1",
|
||||
"type": "ai_yolo",
|
||||
"role": "filter",
|
||||
"enable": true,
|
||||
"model_path": "/models/yolov5s.rknn",
|
||||
"conf": 0.25,
|
||||
"nms": 0.45,
|
||||
"class_filter": []
|
||||
},
|
||||
{
|
||||
"id": "osd_cam1",
|
||||
"type": "osd",
|
||||
"role": "filter",
|
||||
"enable": true,
|
||||
"draw_bbox": true,
|
||||
"draw_text": true,
|
||||
"line_width": 2,
|
||||
"font_scale": 1
|
||||
},
|
||||
{
|
||||
"id": "pub_cam1",
|
||||
"type": "publish",
|
||||
"role": "sink",
|
||||
"enable": true,
|
||||
"codec": "h264",
|
||||
"fps": 25,
|
||||
"gop": 50,
|
||||
"bitrate_kbps": 2000,
|
||||
"use_mpp": true,
|
||||
"use_ffmpeg_mux": true,
|
||||
"outputs": [
|
||||
{ "proto": "rtsp_server", "port": 8554, "path": "/live/cam1" },
|
||||
{ "proto": "hls", "port": 8080, "path": "/hls/cam1", "segment_sec": 2 }
|
||||
]
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
["in_cam1", "pre_cam1"],
|
||||
["pre_cam1", "ai_cam1"],
|
||||
["ai_cam1", "osd_cam1"],
|
||||
["osd_cam1", "pub_cam1"]
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@ -4,6 +4,7 @@ option(RK3588_ENABLE_FFMPEG "Enable FFmpeg-based RTSP input" OFF)
|
||||
option(RK3588_ENABLE_MPP "Enable Rockchip MPP decode/encode" OFF)
|
||||
option(RK3588_ENABLE_ZLMEDIAKIT "Enable embedded ZLMediaKit RTSP server" OFF)
|
||||
option(RK3588_ENABLE_RGA "Enable Rockchip RGA hardware acceleration" OFF)
|
||||
option(RK3588_ENABLE_RKNN "Enable RKNN NPU inference" OFF)
|
||||
|
||||
set(RK_RGA_ROOT "${RK_RKNN_ROOT}/examples/3rdparty/rga/RK3588" CACHE PATH "Path to RGA library")
|
||||
set(RK_RGA_INCLUDE_DIR "${RK_RGA_ROOT}/include" CACHE PATH "RGA include directory")
|
||||
@ -71,6 +72,21 @@ if(RK3588_ENABLE_RGA)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if(RK3588_ENABLE_RKNN)
|
||||
find_library(RK_RKNN_LIB rknnrt
|
||||
HINTS
|
||||
${RKNN_RUNTIME_LIB_DIR}
|
||||
${RK_RKNN_ROOT}/runtime/RK3588/Linux/librknn_api/aarch64
|
||||
NO_DEFAULT_PATH
|
||||
)
|
||||
if(NOT RK_RKNN_LIB)
|
||||
find_library(RK_RKNN_LIB rknnrt)
|
||||
endif()
|
||||
if(NOT RK_RKNN_LIB)
|
||||
message(WARNING "RKNN enabled but librknnrt not found; disable RK3588_ENABLE_RKNN or set RKNN_RUNTIME_LIB_DIR")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
add_library(input_rtsp SHARED input_rtsp/input_rtsp_node.cpp)
|
||||
target_include_directories(input_rtsp PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/third_party)
|
||||
target_link_libraries(input_rtsp PRIVATE project_options Threads::Threads)
|
||||
@ -144,7 +160,32 @@ set_target_properties(preprocess PROPERTIES
|
||||
RUNTIME_OUTPUT_DIRECTORY ${RK_PLUGIN_OUTPUT_DIR}
|
||||
)
|
||||
|
||||
install(TARGETS input_rtsp publish preprocess
|
||||
# ai_yolo plugin (RKNN-based YOLO inference)
|
||||
add_library(ai_yolo SHARED ai_yolo/ai_yolo_node.cpp)
|
||||
target_include_directories(ai_yolo PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/third_party)
|
||||
target_link_libraries(ai_yolo PRIVATE project_options Threads::Threads)
|
||||
if(RK3588_ENABLE_RKNN AND RK_RKNN_LIB)
|
||||
target_compile_definitions(ai_yolo PRIVATE RK3588_ENABLE_RKNN)
|
||||
target_include_directories(ai_yolo PRIVATE ${RKNN_RUNTIME_INCLUDE_DIR})
|
||||
target_link_libraries(ai_yolo PRIVATE ${RK_RKNN_LIB})
|
||||
endif()
|
||||
set_target_properties(ai_yolo PROPERTIES
|
||||
OUTPUT_NAME "ai_yolo"
|
||||
LIBRARY_OUTPUT_DIRECTORY ${RK_PLUGIN_OUTPUT_DIR}
|
||||
RUNTIME_OUTPUT_DIRECTORY ${RK_PLUGIN_OUTPUT_DIR}
|
||||
)
|
||||
|
||||
# osd plugin (on-screen display for detection results)
|
||||
add_library(osd SHARED osd/osd_node.cpp)
|
||||
target_include_directories(osd PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/third_party)
|
||||
target_link_libraries(osd PRIVATE project_options Threads::Threads)
|
||||
set_target_properties(osd PROPERTIES
|
||||
OUTPUT_NAME "osd"
|
||||
LIBRARY_OUTPUT_DIRECTORY ${RK_PLUGIN_OUTPUT_DIR}
|
||||
RUNTIME_OUTPUT_DIRECTORY ${RK_PLUGIN_OUTPUT_DIR}
|
||||
)
|
||||
|
||||
install(TARGETS input_rtsp publish preprocess ai_yolo osd
|
||||
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}/rk3588-media-server/plugins
|
||||
RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}/rk3588-media-server/plugins
|
||||
)
|
||||
|
||||
475
plugins/ai_yolo/ai_yolo_node.cpp
Normal file
475
plugins/ai_yolo/ai_yolo_node.cpp
Normal file
@ -0,0 +1,475 @@
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#include "node.h"
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
#include "rknn_api.h"
|
||||
#endif
|
||||
|
||||
namespace rk3588 {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kObjClassNum = 80;
|
||||
constexpr int kPropBoxSize = 5 + kObjClassNum;
|
||||
constexpr int kMaxDetections = 64;
|
||||
|
||||
const int kAnchor0[6] = {10, 13, 16, 30, 33, 23};
|
||||
const int kAnchor1[6] = {30, 61, 62, 45, 59, 119};
|
||||
const int kAnchor2[6] = {116, 90, 156, 198, 373, 326};
|
||||
|
||||
const char* kCocoLabels[kObjClassNum] = {
|
||||
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
|
||||
"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
|
||||
"dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
|
||||
"umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
|
||||
"kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
|
||||
"bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
|
||||
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
|
||||
"couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
|
||||
"remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
|
||||
"book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
|
||||
};
|
||||
|
||||
inline int Clamp(float val, int min_val, int max_val) {
|
||||
return val > min_val ? (val < max_val ? static_cast<int>(val) : max_val) : min_val;
|
||||
}
|
||||
|
||||
inline int32_t ClipFloat(float val, float min_val, float max_val) {
|
||||
return static_cast<int32_t>(val <= min_val ? min_val : (val >= max_val ? max_val : val));
|
||||
}
|
||||
|
||||
inline int8_t QuantizeF32ToAffine(float f32, int32_t zp, float scale) {
|
||||
float dst_val = (f32 / scale) + zp;
|
||||
return static_cast<int8_t>(ClipFloat(dst_val, -128, 127));
|
||||
}
|
||||
|
||||
inline float DequantizeAffineToF32(int8_t qnt, int32_t zp, float scale) {
|
||||
return (static_cast<float>(qnt) - static_cast<float>(zp)) * scale;
|
||||
}
|
||||
|
||||
float CalculateIoU(float x1_min, float y1_min, float x1_max, float y1_max,
|
||||
float x2_min, float y2_min, float x2_max, float y2_max) {
|
||||
float w = std::fmax(0.f, std::fmin(x1_max, x2_max) - std::fmax(x1_min, x2_min) + 1.0f);
|
||||
float h = std::fmax(0.f, std::fmin(y1_max, y2_max) - std::fmax(y1_min, y2_min) + 1.0f);
|
||||
float inter = w * h;
|
||||
float area1 = (x1_max - x1_min + 1.0f) * (y1_max - y1_min + 1.0f);
|
||||
float area2 = (x2_max - x2_min + 1.0f) * (y2_max - y2_min + 1.0f);
|
||||
float uni = area1 + area2 - inter;
|
||||
return uni <= 0.f ? 0.f : (inter / uni);
|
||||
}
|
||||
|
||||
void QuickSortDescending(std::vector<float>& values, int left, int right, std::vector<int>& indices) {
|
||||
if (left >= right) return;
|
||||
float pivot = values[left];
|
||||
int pivot_idx = indices[left];
|
||||
int low = left, high = right;
|
||||
while (low < high) {
|
||||
while (low < high && values[high] <= pivot) high--;
|
||||
values[low] = values[high];
|
||||
indices[low] = indices[high];
|
||||
while (low < high && values[low] >= pivot) low++;
|
||||
values[high] = values[low];
|
||||
indices[high] = indices[low];
|
||||
}
|
||||
values[low] = pivot;
|
||||
indices[low] = pivot_idx;
|
||||
QuickSortDescending(values, left, low - 1, indices);
|
||||
QuickSortDescending(values, low + 1, right, indices);
|
||||
}
|
||||
|
||||
void NMS(int valid_count, std::vector<float>& boxes, std::vector<int>& class_ids,
|
||||
std::vector<int>& order, int filter_id, float threshold) {
|
||||
for (int i = 0; i < valid_count; ++i) {
|
||||
if (order[i] == -1 || class_ids[i] != filter_id) continue;
|
||||
int n = order[i];
|
||||
for (int j = i + 1; j < valid_count; ++j) {
|
||||
int m = order[j];
|
||||
if (m == -1 || class_ids[j] != filter_id) continue;
|
||||
float x1_min = boxes[n * 4 + 0];
|
||||
float y1_min = boxes[n * 4 + 1];
|
||||
float x1_max = x1_min + boxes[n * 4 + 2];
|
||||
float y1_max = y1_min + boxes[n * 4 + 3];
|
||||
float x2_min = boxes[m * 4 + 0];
|
||||
float y2_min = boxes[m * 4 + 1];
|
||||
float x2_max = x2_min + boxes[m * 4 + 2];
|
||||
float y2_max = y2_min + boxes[m * 4 + 3];
|
||||
if (CalculateIoU(x1_min, y1_min, x1_max, y1_max, x2_min, y2_min, x2_max, y2_max) > threshold) {
|
||||
order[j] = -1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
int ProcessFeatureMap(int8_t* input, const int* anchor, int grid_h, int grid_w,
|
||||
int model_h, int model_w, int stride,
|
||||
std::vector<float>& boxes, std::vector<float>& obj_probs,
|
||||
std::vector<int>& class_ids, float conf_thresh, int32_t zp, float scale) {
|
||||
int valid_count = 0;
|
||||
int grid_len = grid_h * grid_w;
|
||||
int8_t thresh_i8 = QuantizeF32ToAffine(conf_thresh, zp, scale);
|
||||
|
||||
for (int a = 0; a < 3; ++a) {
|
||||
for (int i = 0; i < grid_h; ++i) {
|
||||
for (int j = 0; j < grid_w; ++j) {
|
||||
int8_t box_conf = input[(kPropBoxSize * a + 4) * grid_len + i * grid_w + j];
|
||||
if (box_conf >= thresh_i8) {
|
||||
int offset = (kPropBoxSize * a) * grid_len + i * grid_w + j;
|
||||
int8_t* ptr = input + offset;
|
||||
|
||||
float bx = DequantizeAffineToF32(*ptr, zp, scale) * 2.0f - 0.5f;
|
||||
float by = DequantizeAffineToF32(ptr[grid_len], zp, scale) * 2.0f - 0.5f;
|
||||
float bw = DequantizeAffineToF32(ptr[2 * grid_len], zp, scale) * 2.0f;
|
||||
float bh = DequantizeAffineToF32(ptr[3 * grid_len], zp, scale) * 2.0f;
|
||||
|
||||
bx = (bx + j) * stride;
|
||||
by = (by + i) * stride;
|
||||
bw = bw * bw * anchor[a * 2];
|
||||
bh = bh * bh * anchor[a * 2 + 1];
|
||||
bx -= bw / 2.0f;
|
||||
by -= bh / 2.0f;
|
||||
|
||||
int8_t max_cls_prob = ptr[5 * grid_len];
|
||||
int max_cls_id = 0;
|
||||
for (int k = 1; k < kObjClassNum; ++k) {
|
||||
int8_t prob = ptr[(5 + k) * grid_len];
|
||||
if (prob > max_cls_prob) {
|
||||
max_cls_id = k;
|
||||
max_cls_prob = prob;
|
||||
}
|
||||
}
|
||||
|
||||
if (max_cls_prob > thresh_i8) {
|
||||
float score = DequantizeAffineToF32(max_cls_prob, zp, scale) *
|
||||
DequantizeAffineToF32(box_conf, zp, scale);
|
||||
obj_probs.push_back(score);
|
||||
class_ids.push_back(max_cls_id);
|
||||
boxes.push_back(bx);
|
||||
boxes.push_back(by);
|
||||
boxes.push_back(bw);
|
||||
boxes.push_back(bh);
|
||||
++valid_count;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return valid_count;
|
||||
}
|
||||
#endif
|
||||
|
||||
} // namespace
|
||||
|
||||
class AiYoloNode : public INode {
|
||||
public:
|
||||
std::string Id() const override { return id_; }
|
||||
std::string Type() const override { return "ai_yolo"; }
|
||||
|
||||
bool Init(const SimpleJson& config, const NodeContext& ctx) override {
|
||||
id_ = config.ValueOr<std::string>("id", "ai_yolo");
|
||||
model_path_ = config.ValueOr<std::string>("model_path", "");
|
||||
conf_thresh_ = config.ValueOr<float>("conf", 0.25f);
|
||||
nms_thresh_ = config.ValueOr<float>("nms", 0.45f);
|
||||
model_input_w_ = config.ValueOr<int>("model_w", 640);
|
||||
model_input_h_ = config.ValueOr<int>("model_h", 640);
|
||||
|
||||
if (const SimpleJson* filter = config.Find("class_filter")) {
|
||||
for (const auto& item : filter->AsArray()) {
|
||||
class_filter_.insert(item.AsInt(-1));
|
||||
}
|
||||
}
|
||||
|
||||
input_queue_ = ctx.input_queue;
|
||||
if (!input_queue_) {
|
||||
std::cerr << "[ai_yolo] no input queue for node " << id_ << "\n";
|
||||
return false;
|
||||
}
|
||||
if (ctx.output_queues.empty()) {
|
||||
std::cerr << "[ai_yolo] no output queue for node " << id_ << "\n";
|
||||
return false;
|
||||
}
|
||||
output_queues_ = ctx.output_queues;
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
if (model_path_.empty()) {
|
||||
std::cerr << "[ai_yolo] model_path is required\n";
|
||||
return false;
|
||||
}
|
||||
if (!LoadModel()) {
|
||||
std::cerr << "[ai_yolo] failed to load model: " << model_path_ << "\n";
|
||||
return false;
|
||||
}
|
||||
std::cout << "[ai_yolo] model loaded: " << model_path_ << "\n";
|
||||
#else
|
||||
std::cout << "[ai_yolo] RKNN disabled, will passthrough frames\n";
|
||||
#endif
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Start() override {
|
||||
if (!input_queue_) return false;
|
||||
running_.store(true);
|
||||
worker_ = std::thread(&AiYoloNode::WorkerLoop, this);
|
||||
std::cout << "[ai_yolo] started, conf=" << conf_thresh_ << " nms=" << nms_thresh_ << "\n";
|
||||
return true;
|
||||
}
|
||||
|
||||
void Stop() override {
|
||||
running_.store(false);
|
||||
if (input_queue_) input_queue_->Stop();
|
||||
for (auto& q : output_queues_) q->Stop();
|
||||
if (worker_.joinable()) worker_.join();
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
if (rknn_ctx_) {
|
||||
rknn_destroy(rknn_ctx_);
|
||||
rknn_ctx_ = 0;
|
||||
}
|
||||
#endif
|
||||
std::cout << "[ai_yolo] stopped\n";
|
||||
}
|
||||
|
||||
private:
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
bool LoadModel() {
|
||||
std::ifstream file(model_path_, std::ios::binary | std::ios::ate);
|
||||
if (!file.is_open()) return false;
|
||||
|
||||
size_t model_size = file.tellg();
|
||||
file.seekg(0, std::ios::beg);
|
||||
model_data_.resize(model_size);
|
||||
if (!file.read(reinterpret_cast<char*>(model_data_.data()), model_size)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
int ret = rknn_init(&rknn_ctx_, model_data_.data(), model_size, 0, nullptr);
|
||||
if (ret < 0) {
|
||||
std::cerr << "[ai_yolo] rknn_init failed: " << ret << "\n";
|
||||
return false;
|
||||
}
|
||||
|
||||
rknn_input_output_num io_num;
|
||||
ret = rknn_query(rknn_ctx_, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
|
||||
if (ret < 0) {
|
||||
std::cerr << "[ai_yolo] rknn_query IO num failed\n";
|
||||
return false;
|
||||
}
|
||||
n_input_ = io_num.n_input;
|
||||
n_output_ = io_num.n_output;
|
||||
|
||||
input_attrs_.resize(n_input_);
|
||||
for (uint32_t i = 0; i < n_input_; ++i) {
|
||||
input_attrs_[i].index = i;
|
||||
rknn_query(rknn_ctx_, RKNN_QUERY_INPUT_ATTR, &input_attrs_[i], sizeof(rknn_tensor_attr));
|
||||
}
|
||||
|
||||
output_attrs_.resize(n_output_);
|
||||
for (uint32_t i = 0; i < n_output_; ++i) {
|
||||
output_attrs_[i].index = i;
|
||||
rknn_query(rknn_ctx_, RKNN_QUERY_OUTPUT_ATTR, &output_attrs_[i], sizeof(rknn_tensor_attr));
|
||||
}
|
||||
|
||||
if (input_attrs_[0].fmt == RKNN_TENSOR_NCHW) {
|
||||
model_input_h_ = input_attrs_[0].dims[2];
|
||||
model_input_w_ = input_attrs_[0].dims[3];
|
||||
} else {
|
||||
model_input_h_ = input_attrs_[0].dims[1];
|
||||
model_input_w_ = input_attrs_[0].dims[2];
|
||||
}
|
||||
|
||||
std::cout << "[ai_yolo] model input: " << model_input_w_ << "x" << model_input_h_
|
||||
<< ", outputs: " << n_output_ << "\n";
|
||||
return true;
|
||||
}
|
||||
#endif
|
||||
|
||||
void PushToDownstream(FramePtr frame) {
|
||||
for (auto& q : output_queues_) {
|
||||
q->Push(frame);
|
||||
}
|
||||
}
|
||||
|
||||
void WorkerLoop() {
|
||||
using namespace std::chrono;
|
||||
FramePtr frame;
|
||||
|
||||
while (running_.load()) {
|
||||
if (!input_queue_->Pop(frame, milliseconds(200))) continue;
|
||||
if (!frame) continue;
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
RunInference(frame);
|
||||
#endif
|
||||
PushToDownstream(frame);
|
||||
++processed_;
|
||||
|
||||
if (processed_ % 100 == 0) {
|
||||
std::cout << "[ai_yolo] processed " << processed_ << " frames\n";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
void RunInference(FramePtr frame) {
|
||||
if (!frame->data || frame->data_size == 0) return;
|
||||
|
||||
bool is_rgb = (frame->format == PixelFormat::RGB || frame->format == PixelFormat::BGR);
|
||||
if (!is_rgb) {
|
||||
std::cerr << "[ai_yolo] input must be RGB/BGR, got other format\n";
|
||||
return;
|
||||
}
|
||||
|
||||
rknn_input inputs[1];
|
||||
memset(inputs, 0, sizeof(inputs));
|
||||
inputs[0].index = 0;
|
||||
inputs[0].type = RKNN_TENSOR_UINT8;
|
||||
inputs[0].size = frame->width * frame->height * 3;
|
||||
inputs[0].fmt = RKNN_TENSOR_NHWC;
|
||||
inputs[0].buf = frame->data;
|
||||
inputs[0].pass_through = 0;
|
||||
|
||||
int ret = rknn_inputs_set(rknn_ctx_, n_input_, inputs);
|
||||
if (ret < 0) {
|
||||
std::cerr << "[ai_yolo] rknn_inputs_set failed: " << ret << "\n";
|
||||
return;
|
||||
}
|
||||
|
||||
ret = rknn_run(rknn_ctx_, nullptr);
|
||||
if (ret < 0) {
|
||||
std::cerr << "[ai_yolo] rknn_run failed: " << ret << "\n";
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<rknn_output> outputs(n_output_);
|
||||
memset(outputs.data(), 0, sizeof(rknn_output) * n_output_);
|
||||
for (uint32_t i = 0; i < n_output_; ++i) {
|
||||
outputs[i].want_float = 0;
|
||||
}
|
||||
|
||||
ret = rknn_outputs_get(rknn_ctx_, n_output_, outputs.data(), nullptr);
|
||||
if (ret < 0) {
|
||||
std::cerr << "[ai_yolo] rknn_outputs_get failed: " << ret << "\n";
|
||||
return;
|
||||
}
|
||||
|
||||
PostProcess(outputs, frame);
|
||||
rknn_outputs_release(rknn_ctx_, n_output_, outputs.data());
|
||||
}
|
||||
|
||||
void PostProcess(std::vector<rknn_output>& outputs, FramePtr frame) {
|
||||
if (n_output_ < 3) return;
|
||||
|
||||
std::vector<float> boxes;
|
||||
std::vector<float> obj_probs;
|
||||
std::vector<int> class_ids;
|
||||
|
||||
std::vector<int32_t> zps;
|
||||
std::vector<float> scales;
|
||||
for (uint32_t i = 0; i < n_output_; ++i) {
|
||||
zps.push_back(output_attrs_[i].zp);
|
||||
scales.push_back(output_attrs_[i].scale);
|
||||
}
|
||||
|
||||
int stride0 = 8, stride1 = 16, stride2 = 32;
|
||||
int grid_h0 = model_input_h_ / stride0, grid_w0 = model_input_w_ / stride0;
|
||||
int grid_h1 = model_input_h_ / stride1, grid_w1 = model_input_w_ / stride1;
|
||||
int grid_h2 = model_input_h_ / stride2, grid_w2 = model_input_w_ / stride2;
|
||||
|
||||
int cnt0 = ProcessFeatureMap(reinterpret_cast<int8_t*>(outputs[0].buf), kAnchor0,
|
||||
grid_h0, grid_w0, model_input_h_, model_input_w_, stride0,
|
||||
boxes, obj_probs, class_ids, conf_thresh_, zps[0], scales[0]);
|
||||
int cnt1 = ProcessFeatureMap(reinterpret_cast<int8_t*>(outputs[1].buf), kAnchor1,
|
||||
grid_h1, grid_w1, model_input_h_, model_input_w_, stride1,
|
||||
boxes, obj_probs, class_ids, conf_thresh_, zps[1], scales[1]);
|
||||
int cnt2 = ProcessFeatureMap(reinterpret_cast<int8_t*>(outputs[2].buf), kAnchor2,
|
||||
grid_h2, grid_w2, model_input_h_, model_input_w_, stride2,
|
||||
boxes, obj_probs, class_ids, conf_thresh_, zps[2], scales[2]);
|
||||
|
||||
int valid_count = cnt0 + cnt1 + cnt2;
|
||||
if (valid_count <= 0) return;
|
||||
|
||||
std::vector<int> indices(valid_count);
|
||||
for (int i = 0; i < valid_count; ++i) indices[i] = i;
|
||||
|
||||
QuickSortDescending(obj_probs, 0, valid_count - 1, indices);
|
||||
|
||||
std::set<int> class_set(class_ids.begin(), class_ids.end());
|
||||
for (int c : class_set) {
|
||||
NMS(valid_count, boxes, class_ids, indices, c, nms_thresh_);
|
||||
}
|
||||
|
||||
float scale_w = static_cast<float>(model_input_w_) / frame->width;
|
||||
float scale_h = static_cast<float>(model_input_h_) / frame->height;
|
||||
|
||||
auto det_result = std::make_shared<DetectionResult>();
|
||||
det_result->img_w = frame->width;
|
||||
det_result->img_h = frame->height;
|
||||
det_result->model_name = "yolov5";
|
||||
|
||||
for (int i = 0; i < valid_count && det_result->items.size() < kMaxDetections; ++i) {
|
||||
if (indices[i] == -1) continue;
|
||||
int n = indices[i];
|
||||
int cls_id = class_ids[n];
|
||||
|
||||
if (!class_filter_.empty() && class_filter_.find(cls_id) == class_filter_.end()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
float x1 = boxes[n * 4 + 0];
|
||||
float y1 = boxes[n * 4 + 1];
|
||||
float w = boxes[n * 4 + 2];
|
||||
float h = boxes[n * 4 + 3];
|
||||
|
||||
Detection det;
|
||||
det.cls_id = cls_id;
|
||||
det.score = obj_probs[i];
|
||||
det.bbox.x = Clamp(x1 / scale_w, 0, frame->width);
|
||||
det.bbox.y = Clamp(y1 / scale_h, 0, frame->height);
|
||||
det.bbox.w = Clamp(w / scale_w, 0, frame->width - det.bbox.x);
|
||||
det.bbox.h = Clamp(h / scale_h, 0, frame->height - det.bbox.y);
|
||||
det.track_id = -1;
|
||||
|
||||
det_result->items.push_back(det);
|
||||
}
|
||||
|
||||
frame->det = det_result;
|
||||
}
|
||||
#endif
|
||||
|
||||
std::string id_;
|
||||
std::string model_path_;
|
||||
float conf_thresh_ = 0.25f;
|
||||
float nms_thresh_ = 0.45f;
|
||||
int model_input_w_ = 640;
|
||||
int model_input_h_ = 640;
|
||||
std::set<int> class_filter_;
|
||||
|
||||
std::atomic<bool> running_{false};
|
||||
std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
|
||||
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
|
||||
std::thread worker_;
|
||||
uint64_t processed_ = 0;
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
rknn_context rknn_ctx_ = 0;
|
||||
std::vector<uint8_t> model_data_;
|
||||
uint32_t n_input_ = 0;
|
||||
uint32_t n_output_ = 0;
|
||||
std::vector<rknn_tensor_attr> input_attrs_;
|
||||
std::vector<rknn_tensor_attr> output_attrs_;
|
||||
#endif
|
||||
};
|
||||
|
||||
REGISTER_NODE(AiYoloNode, "ai_yolo");
|
||||
|
||||
} // namespace rk3588
|
||||
397
plugins/osd/osd_node.cpp
Normal file
397
plugins/osd/osd_node.cpp
Normal file
@ -0,0 +1,397 @@
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#include "node.h"
|
||||
|
||||
namespace rk3588 {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kObjClassNum = 80;
|
||||
|
||||
const char* kCocoLabels[kObjClassNum] = {
|
||||
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
|
||||
"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
|
||||
"dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
|
||||
"umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
|
||||
"kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
|
||||
"bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
|
||||
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
|
||||
"couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
|
||||
"remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
|
||||
"book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
|
||||
};
|
||||
|
||||
struct Color {
|
||||
uint8_t r, g, b;
|
||||
};
|
||||
|
||||
const Color kClassColors[] = {
|
||||
{255, 0, 0}, {0, 255, 0}, {0, 0, 255}, {255, 255, 0}, {255, 0, 255},
|
||||
{0, 255, 255}, {128, 0, 0}, {0, 128, 0}, {0, 0, 128}, {128, 128, 0},
|
||||
{128, 0, 128}, {0, 128, 128}, {255, 128, 0}, {255, 0, 128}, {128, 255, 0},
|
||||
{0, 255, 128}, {128, 0, 255}, {0, 128, 255}, {255, 128, 128}, {128, 255, 128}
|
||||
};
|
||||
|
||||
inline Color GetClassColor(int cls_id) {
|
||||
return kClassColors[cls_id % 20];
|
||||
}
|
||||
|
||||
inline const char* GetClassName(int cls_id) {
|
||||
if (cls_id >= 0 && cls_id < kObjClassNum) {
|
||||
return kCocoLabels[cls_id];
|
||||
}
|
||||
return "unknown";
|
||||
}
|
||||
|
||||
inline int Clamp(int val, int min_val, int max_val) {
|
||||
return val < min_val ? min_val : (val > max_val ? max_val : val);
|
||||
}
|
||||
|
||||
void DrawHLine(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
||||
int x1, int x2, int y, int thickness, const Color& color) {
|
||||
x1 = Clamp(x1, 0, w - 1);
|
||||
x2 = Clamp(x2, 0, w - 1);
|
||||
if (x1 > x2) std::swap(x1, x2);
|
||||
|
||||
for (int t = 0; t < thickness; ++t) {
|
||||
int cy = y + t;
|
||||
if (cy < 0 || cy >= h) continue;
|
||||
|
||||
if (fmt == PixelFormat::RGB) {
|
||||
for (int x = x1; x <= x2; ++x) {
|
||||
int idx = (cy * stride) + x * 3;
|
||||
data[idx] = color.r;
|
||||
data[idx + 1] = color.g;
|
||||
data[idx + 2] = color.b;
|
||||
}
|
||||
} else if (fmt == PixelFormat::BGR) {
|
||||
for (int x = x1; x <= x2; ++x) {
|
||||
int idx = (cy * stride) + x * 3;
|
||||
data[idx] = color.b;
|
||||
data[idx + 1] = color.g;
|
||||
data[idx + 2] = color.r;
|
||||
}
|
||||
} else if (fmt == PixelFormat::NV12) {
|
||||
uint8_t Y = static_cast<uint8_t>(0.299f * color.r + 0.587f * color.g + 0.114f * color.b);
|
||||
for (int x = x1; x <= x2; ++x) {
|
||||
data[cy * w + x] = Y;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void DrawVLine(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
||||
int x, int y1, int y2, int thickness, const Color& color) {
|
||||
y1 = Clamp(y1, 0, h - 1);
|
||||
y2 = Clamp(y2, 0, h - 1);
|
||||
if (y1 > y2) std::swap(y1, y2);
|
||||
|
||||
for (int t = 0; t < thickness; ++t) {
|
||||
int cx = x + t;
|
||||
if (cx < 0 || cx >= w) continue;
|
||||
|
||||
if (fmt == PixelFormat::RGB) {
|
||||
for (int y = y1; y <= y2; ++y) {
|
||||
int idx = (y * stride) + cx * 3;
|
||||
data[idx] = color.r;
|
||||
data[idx + 1] = color.g;
|
||||
data[idx + 2] = color.b;
|
||||
}
|
||||
} else if (fmt == PixelFormat::BGR) {
|
||||
for (int y = y1; y <= y2; ++y) {
|
||||
int idx = (y * stride) + cx * 3;
|
||||
data[idx] = color.b;
|
||||
data[idx + 1] = color.g;
|
||||
data[idx + 2] = color.r;
|
||||
}
|
||||
} else if (fmt == PixelFormat::NV12) {
|
||||
uint8_t Y = static_cast<uint8_t>(0.299f * color.r + 0.587f * color.g + 0.114f * color.b);
|
||||
for (int y = y1; y <= y2; ++y) {
|
||||
data[y * w + cx] = Y;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void DrawRect(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
||||
int x1, int y1, int x2, int y2, int thickness, const Color& color) {
|
||||
DrawHLine(data, w, h, stride, fmt, x1, x2, y1, thickness, color);
|
||||
DrawHLine(data, w, h, stride, fmt, x1, x2, y2 - thickness + 1, thickness, color);
|
||||
DrawVLine(data, w, h, stride, fmt, x1, y1, y2, thickness, color);
|
||||
DrawVLine(data, w, h, stride, fmt, x2 - thickness + 1, y1, y2, thickness, color);
|
||||
}
|
||||
|
||||
const uint8_t kFont5x7[96][7] = {
|
||||
{0x00,0x00,0x00,0x00,0x00,0x00,0x00}, // ' '
|
||||
{0x04,0x04,0x04,0x04,0x00,0x00,0x04}, // '!'
|
||||
{0x0A,0x0A,0x0A,0x00,0x00,0x00,0x00}, // '"'
|
||||
{0x0A,0x0A,0x1F,0x0A,0x1F,0x0A,0x0A}, // '#'
|
||||
{0x04,0x0F,0x14,0x0E,0x05,0x1E,0x04}, // '$'
|
||||
{0x18,0x19,0x02,0x04,0x08,0x13,0x03}, // '%'
|
||||
{0x0C,0x12,0x14,0x08,0x15,0x12,0x0D}, // '&'
|
||||
{0x0C,0x04,0x08,0x00,0x00,0x00,0x00}, // '''
|
||||
{0x02,0x04,0x08,0x08,0x08,0x04,0x02}, // '('
|
||||
{0x08,0x04,0x02,0x02,0x02,0x04,0x08}, // ')'
|
||||
{0x00,0x04,0x15,0x0E,0x15,0x04,0x00}, // '*'
|
||||
{0x00,0x04,0x04,0x1F,0x04,0x04,0x00}, // '+'
|
||||
{0x00,0x00,0x00,0x00,0x0C,0x04,0x08}, // ','
|
||||
{0x00,0x00,0x00,0x1F,0x00,0x00,0x00}, // '-'
|
||||
{0x00,0x00,0x00,0x00,0x00,0x0C,0x0C}, // '.'
|
||||
{0x00,0x01,0x02,0x04,0x08,0x10,0x00}, // '/'
|
||||
{0x0E,0x11,0x13,0x15,0x19,0x11,0x0E}, // '0'
|
||||
{0x04,0x0C,0x04,0x04,0x04,0x04,0x0E}, // '1'
|
||||
{0x0E,0x11,0x01,0x02,0x04,0x08,0x1F}, // '2'
|
||||
{0x1F,0x02,0x04,0x02,0x01,0x11,0x0E}, // '3'
|
||||
{0x02,0x06,0x0A,0x12,0x1F,0x02,0x02}, // '4'
|
||||
{0x1F,0x10,0x1E,0x01,0x01,0x11,0x0E}, // '5'
|
||||
{0x06,0x08,0x10,0x1E,0x11,0x11,0x0E}, // '6'
|
||||
{0x1F,0x01,0x02,0x04,0x08,0x08,0x08}, // '7'
|
||||
{0x0E,0x11,0x11,0x0E,0x11,0x11,0x0E}, // '8'
|
||||
{0x0E,0x11,0x11,0x0F,0x01,0x02,0x0C}, // '9'
|
||||
{0x00,0x0C,0x0C,0x00,0x0C,0x0C,0x00}, // ':'
|
||||
{0x00,0x0C,0x0C,0x00,0x0C,0x04,0x08}, // ';'
|
||||
{0x02,0x04,0x08,0x10,0x08,0x04,0x02}, // '<'
|
||||
{0x00,0x00,0x1F,0x00,0x1F,0x00,0x00}, // '='
|
||||
{0x08,0x04,0x02,0x01,0x02,0x04,0x08}, // '>'
|
||||
{0x0E,0x11,0x01,0x02,0x04,0x00,0x04}, // '?'
|
||||
{0x0E,0x11,0x17,0x15,0x17,0x10,0x0E}, // '@'
|
||||
{0x0E,0x11,0x11,0x1F,0x11,0x11,0x11}, // 'A'
|
||||
{0x1E,0x11,0x11,0x1E,0x11,0x11,0x1E}, // 'B'
|
||||
{0x0E,0x11,0x10,0x10,0x10,0x11,0x0E}, // 'C'
|
||||
{0x1C,0x12,0x11,0x11,0x11,0x12,0x1C}, // 'D'
|
||||
{0x1F,0x10,0x10,0x1E,0x10,0x10,0x1F}, // 'E'
|
||||
{0x1F,0x10,0x10,0x1E,0x10,0x10,0x10}, // 'F'
|
||||
{0x0E,0x11,0x10,0x17,0x11,0x11,0x0F}, // 'G'
|
||||
{0x11,0x11,0x11,0x1F,0x11,0x11,0x11}, // 'H'
|
||||
{0x0E,0x04,0x04,0x04,0x04,0x04,0x0E}, // 'I'
|
||||
{0x07,0x02,0x02,0x02,0x02,0x12,0x0C}, // 'J'
|
||||
{0x11,0x12,0x14,0x18,0x14,0x12,0x11}, // 'K'
|
||||
{0x10,0x10,0x10,0x10,0x10,0x10,0x1F}, // 'L'
|
||||
{0x11,0x1B,0x15,0x15,0x11,0x11,0x11}, // 'M'
|
||||
{0x11,0x11,0x19,0x15,0x13,0x11,0x11}, // 'N'
|
||||
{0x0E,0x11,0x11,0x11,0x11,0x11,0x0E}, // 'O'
|
||||
{0x1E,0x11,0x11,0x1E,0x10,0x10,0x10}, // 'P'
|
||||
{0x0E,0x11,0x11,0x11,0x15,0x12,0x0D}, // 'Q'
|
||||
{0x1E,0x11,0x11,0x1E,0x14,0x12,0x11}, // 'R'
|
||||
{0x0F,0x10,0x10,0x0E,0x01,0x01,0x1E}, // 'S'
|
||||
{0x1F,0x04,0x04,0x04,0x04,0x04,0x04}, // 'T'
|
||||
{0x11,0x11,0x11,0x11,0x11,0x11,0x0E}, // 'U'
|
||||
{0x11,0x11,0x11,0x11,0x11,0x0A,0x04}, // 'V'
|
||||
{0x11,0x11,0x11,0x15,0x15,0x15,0x0A}, // 'W'
|
||||
{0x11,0x11,0x0A,0x04,0x0A,0x11,0x11}, // 'X'
|
||||
{0x11,0x11,0x11,0x0A,0x04,0x04,0x04}, // 'Y'
|
||||
{0x1F,0x01,0x02,0x04,0x08,0x10,0x1F}, // 'Z'
|
||||
{0x0E,0x08,0x08,0x08,0x08,0x08,0x0E}, // '['
|
||||
{0x00,0x10,0x08,0x04,0x02,0x01,0x00}, // '\'
|
||||
{0x0E,0x02,0x02,0x02,0x02,0x02,0x0E}, // ']'
|
||||
{0x04,0x0A,0x11,0x00,0x00,0x00,0x00}, // '^'
|
||||
{0x00,0x00,0x00,0x00,0x00,0x00,0x1F}, // '_'
|
||||
{0x08,0x04,0x02,0x00,0x00,0x00,0x00}, // '`'
|
||||
{0x00,0x00,0x0E,0x01,0x0F,0x11,0x0F}, // 'a'
|
||||
{0x10,0x10,0x16,0x19,0x11,0x11,0x1E}, // 'b'
|
||||
{0x00,0x00,0x0E,0x10,0x10,0x11,0x0E}, // 'c'
|
||||
{0x01,0x01,0x0D,0x13,0x11,0x11,0x0F}, // 'd'
|
||||
{0x00,0x00,0x0E,0x11,0x1F,0x10,0x0E}, // 'e'
|
||||
{0x06,0x09,0x08,0x1C,0x08,0x08,0x08}, // 'f'
|
||||
{0x00,0x0F,0x11,0x11,0x0F,0x01,0x0E}, // 'g'
|
||||
{0x10,0x10,0x16,0x19,0x11,0x11,0x11}, // 'h'
|
||||
{0x04,0x00,0x0C,0x04,0x04,0x04,0x0E}, // 'i'
|
||||
{0x02,0x00,0x06,0x02,0x02,0x12,0x0C}, // 'j'
|
||||
{0x10,0x10,0x12,0x14,0x18,0x14,0x12}, // 'k'
|
||||
{0x0C,0x04,0x04,0x04,0x04,0x04,0x0E}, // 'l'
|
||||
{0x00,0x00,0x1A,0x15,0x15,0x11,0x11}, // 'm'
|
||||
{0x00,0x00,0x16,0x19,0x11,0x11,0x11}, // 'n'
|
||||
{0x00,0x00,0x0E,0x11,0x11,0x11,0x0E}, // 'o'
|
||||
{0x00,0x00,0x1E,0x11,0x1E,0x10,0x10}, // 'p'
|
||||
{0x00,0x00,0x0D,0x13,0x0F,0x01,0x01}, // 'q'
|
||||
{0x00,0x00,0x16,0x19,0x10,0x10,0x10}, // 'r'
|
||||
{0x00,0x00,0x0E,0x10,0x0E,0x01,0x1E}, // 's'
|
||||
{0x08,0x08,0x1C,0x08,0x08,0x09,0x06}, // 't'
|
||||
{0x00,0x00,0x11,0x11,0x11,0x13,0x0D}, // 'u'
|
||||
{0x00,0x00,0x11,0x11,0x11,0x0A,0x04}, // 'v'
|
||||
{0x00,0x00,0x11,0x11,0x15,0x15,0x0A}, // 'w'
|
||||
{0x00,0x00,0x11,0x0A,0x04,0x0A,0x11}, // 'x'
|
||||
{0x00,0x00,0x11,0x11,0x0F,0x01,0x0E}, // 'y'
|
||||
{0x00,0x00,0x1F,0x02,0x04,0x08,0x1F}, // 'z'
|
||||
{0x02,0x04,0x04,0x08,0x04,0x04,0x02}, // '{'
|
||||
{0x04,0x04,0x04,0x04,0x04,0x04,0x04}, // '|'
|
||||
{0x08,0x04,0x04,0x02,0x04,0x04,0x08}, // '}'
|
||||
{0x00,0x00,0x08,0x15,0x02,0x00,0x00}, // '~'
|
||||
{0x00,0x00,0x00,0x00,0x00,0x00,0x00}, // DEL
|
||||
};
|
||||
|
||||
void DrawChar(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
||||
int x, int y, char c, int scale, const Color& color) {
|
||||
if (c < 32 || c > 127) c = ' ';
|
||||
int idx = c - 32;
|
||||
const uint8_t* glyph = kFont5x7[idx];
|
||||
|
||||
for (int row = 0; row < 7; ++row) {
|
||||
for (int col = 0; col < 5; ++col) {
|
||||
if (glyph[row] & (1 << (4 - col))) {
|
||||
for (int sy = 0; sy < scale; ++sy) {
|
||||
for (int sx = 0; sx < scale; ++sx) {
|
||||
int px = x + col * scale + sx;
|
||||
int py = y + row * scale + sy;
|
||||
if (px < 0 || px >= w || py < 0 || py >= h) continue;
|
||||
|
||||
if (fmt == PixelFormat::RGB) {
|
||||
int i = py * stride + px * 3;
|
||||
data[i] = color.r;
|
||||
data[i + 1] = color.g;
|
||||
data[i + 2] = color.b;
|
||||
} else if (fmt == PixelFormat::BGR) {
|
||||
int i = py * stride + px * 3;
|
||||
data[i] = color.b;
|
||||
data[i + 1] = color.g;
|
||||
data[i + 2] = color.r;
|
||||
} else if (fmt == PixelFormat::NV12) {
|
||||
uint8_t Y = static_cast<uint8_t>(0.299f * color.r + 0.587f * color.g + 0.114f * color.b);
|
||||
data[py * w + px] = Y;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void DrawText(uint8_t* data, int w, int h, int stride, PixelFormat fmt,
|
||||
int x, int y, const char* text, int scale, const Color& color) {
|
||||
int cx = x;
|
||||
while (*text) {
|
||||
DrawChar(data, w, h, stride, fmt, cx, y, *text, scale, color);
|
||||
cx += 6 * scale;
|
||||
++text;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
class OsdNode : public INode {
|
||||
public:
|
||||
std::string Id() const override { return id_; }
|
||||
std::string Type() const override { return "osd"; }
|
||||
|
||||
bool Init(const SimpleJson& config, const NodeContext& ctx) override {
|
||||
id_ = config.ValueOr<std::string>("id", "osd");
|
||||
draw_bbox_ = config.ValueOr<bool>("draw_bbox", true);
|
||||
draw_text_ = config.ValueOr<bool>("draw_text", true);
|
||||
line_width_ = config.ValueOr<int>("line_width", 2);
|
||||
font_scale_ = config.ValueOr<int>("font_scale", 1);
|
||||
|
||||
input_queue_ = ctx.input_queue;
|
||||
if (!input_queue_) {
|
||||
std::cerr << "[osd] no input queue for node " << id_ << "\n";
|
||||
return false;
|
||||
}
|
||||
if (ctx.output_queues.empty()) {
|
||||
std::cerr << "[osd] no output queue for node " << id_ << "\n";
|
||||
return false;
|
||||
}
|
||||
output_queues_ = ctx.output_queues;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool Start() override {
|
||||
if (!input_queue_) return false;
|
||||
running_.store(true);
|
||||
worker_ = std::thread(&OsdNode::WorkerLoop, this);
|
||||
std::cout << "[osd] started, draw_bbox=" << draw_bbox_ << " draw_text=" << draw_text_ << "\n";
|
||||
return true;
|
||||
}
|
||||
|
||||
void Stop() override {
|
||||
running_.store(false);
|
||||
if (input_queue_) input_queue_->Stop();
|
||||
for (auto& q : output_queues_) q->Stop();
|
||||
if (worker_.joinable()) worker_.join();
|
||||
std::cout << "[osd] stopped\n";
|
||||
}
|
||||
|
||||
private:
|
||||
void PushToDownstream(FramePtr frame) {
|
||||
for (auto& q : output_queues_) {
|
||||
q->Push(frame);
|
||||
}
|
||||
}
|
||||
|
||||
void WorkerLoop() {
|
||||
using namespace std::chrono;
|
||||
FramePtr frame;
|
||||
|
||||
while (running_.load()) {
|
||||
if (!input_queue_->Pop(frame, milliseconds(200))) continue;
|
||||
if (!frame) continue;
|
||||
|
||||
if (frame->det && frame->data) {
|
||||
DrawDetections(frame);
|
||||
}
|
||||
|
||||
PushToDownstream(frame);
|
||||
++processed_;
|
||||
|
||||
if (processed_ % 100 == 0) {
|
||||
std::cout << "[osd] processed " << processed_ << " frames\n";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void DrawDetections(FramePtr frame) {
|
||||
if (!frame->det || frame->det->items.empty()) return;
|
||||
|
||||
int w = frame->width;
|
||||
int h = frame->height;
|
||||
int stride = frame->stride > 0 ? frame->stride : w * 3;
|
||||
uint8_t* data = frame->data;
|
||||
PixelFormat fmt = frame->format;
|
||||
|
||||
bool supported = (fmt == PixelFormat::RGB || fmt == PixelFormat::BGR || fmt == PixelFormat::NV12);
|
||||
if (!supported) {
|
||||
return;
|
||||
}
|
||||
|
||||
for (const auto& det : frame->det->items) {
|
||||
int x1 = static_cast<int>(det.bbox.x);
|
||||
int y1 = static_cast<int>(det.bbox.y);
|
||||
int x2 = static_cast<int>(det.bbox.x + det.bbox.w);
|
||||
int y2 = static_cast<int>(det.bbox.y + det.bbox.h);
|
||||
|
||||
Color color = GetClassColor(det.cls_id);
|
||||
|
||||
if (draw_bbox_) {
|
||||
DrawRect(data, w, h, stride, fmt, x1, y1, x2, y2, line_width_, color);
|
||||
}
|
||||
|
||||
if (draw_text_) {
|
||||
char label[64];
|
||||
snprintf(label, sizeof(label), "%s %.0f%%", GetClassName(det.cls_id), det.score * 100);
|
||||
int text_y = y1 - 8 * font_scale_;
|
||||
if (text_y < 0) text_y = y1 + 2;
|
||||
DrawText(data, w, h, stride, fmt, x1, text_y, label, font_scale_, color);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::string id_;
|
||||
bool draw_bbox_ = true;
|
||||
bool draw_text_ = true;
|
||||
int line_width_ = 2;
|
||||
int font_scale_ = 1;
|
||||
|
||||
std::atomic<bool> running_{false};
|
||||
std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
|
||||
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
|
||||
std::thread worker_;
|
||||
uint64_t processed_ = 0;
|
||||
};
|
||||
|
||||
REGISTER_NODE(OsdNode, "osd");
|
||||
|
||||
} // namespace rk3588
|
||||
Loading…
Reference in New Issue
Block a user