diff --git a/configs/sample_ai_pipeline.json b/configs/sample_ai_pipeline.json new file mode 100644 index 0000000..dbb6a5f --- /dev/null +++ b/configs/sample_ai_pipeline.json @@ -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"] + ] + } + ] +} diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt index afb9b37..c91eac2 100644 --- a/plugins/CMakeLists.txt +++ b/plugins/CMakeLists.txt @@ -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 ) diff --git a/plugins/ai_yolo/ai_yolo_node.cpp b/plugins/ai_yolo/ai_yolo_node.cpp new file mode 100644 index 0000000..7f0321f --- /dev/null +++ b/plugins/ai_yolo/ai_yolo_node.cpp @@ -0,0 +1,475 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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(val) : max_val) : min_val; +} + +inline int32_t ClipFloat(float val, float min_val, float max_val) { + return static_cast(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(ClipFloat(dst_val, -128, 127)); +} + +inline float DequantizeAffineToF32(int8_t qnt, int32_t zp, float scale) { + return (static_cast(qnt) - static_cast(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& values, int left, int right, std::vector& 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& boxes, std::vector& class_ids, + std::vector& 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& boxes, std::vector& obj_probs, + std::vector& 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("id", "ai_yolo"); + model_path_ = config.ValueOr("model_path", ""); + conf_thresh_ = config.ValueOr("conf", 0.25f); + nms_thresh_ = config.ValueOr("nms", 0.45f); + model_input_w_ = config.ValueOr("model_w", 640); + model_input_h_ = config.ValueOr("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(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 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& outputs, FramePtr frame) { + if (n_output_ < 3) return; + + std::vector boxes; + std::vector obj_probs; + std::vector class_ids; + + std::vector zps; + std::vector 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(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(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(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 indices(valid_count); + for (int i = 0; i < valid_count; ++i) indices[i] = i; + + QuickSortDescending(obj_probs, 0, valid_count - 1, indices); + + std::set 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(model_input_w_) / frame->width; + float scale_h = static_cast(model_input_h_) / frame->height; + + auto det_result = std::make_shared(); + 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 class_filter_; + + std::atomic running_{false}; + std::shared_ptr> input_queue_; + std::vector>> output_queues_; + std::thread worker_; + uint64_t processed_ = 0; + +#if defined(RK3588_ENABLE_RKNN) + rknn_context rknn_ctx_ = 0; + std::vector model_data_; + uint32_t n_input_ = 0; + uint32_t n_output_ = 0; + std::vector input_attrs_; + std::vector output_attrs_; +#endif +}; + +REGISTER_NODE(AiYoloNode, "ai_yolo"); + +} // namespace rk3588 diff --git a/plugins/osd/osd_node.cpp b/plugins/osd/osd_node.cpp new file mode 100644 index 0000000..2da95e8 --- /dev/null +++ b/plugins/osd/osd_node.cpp @@ -0,0 +1,397 @@ +#include +#include +#include +#include +#include +#include +#include + +#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(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(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(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("id", "osd"); + draw_bbox_ = config.ValueOr("draw_bbox", true); + draw_text_ = config.ValueOr("draw_text", true); + line_width_ = config.ValueOr("line_width", 2); + font_scale_ = config.ValueOr("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(det.bbox.x); + int y1 = static_cast(det.bbox.y); + int x2 = static_cast(det.bbox.x + det.bbox.w); + int y2 = static_cast(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 running_{false}; + std::shared_ptr> input_queue_; + std::vector>> output_queues_; + std::thread worker_; + uint64_t processed_ = 0; +}; + +REGISTER_NODE(OsdNode, "osd"); + +} // namespace rk3588