diff --git a/plugins/CMakeLists.txt b/plugins/CMakeLists.txt index 747be42..04e2a3e 100644 --- a/plugins/CMakeLists.txt +++ b/plugins/CMakeLists.txt @@ -221,6 +221,16 @@ set_target_properties(ai_yolo PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${RK_PLUGIN_OUTPUT_DIR} ) +# det_post plugin (post-processing pipeline for detection results) +add_library(det_post SHARED det_post/det_post_node.cpp) +target_include_directories(det_post PRIVATE ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/third_party) +target_link_libraries(det_post PRIVATE project_options Threads::Threads) +set_target_properties(det_post PROPERTIES + OUTPUT_NAME "det_post" + 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) diff --git a/plugins/det_post/det_post_node.cpp b/plugins/det_post/det_post_node.cpp new file mode 100644 index 0000000..83f77c7 --- /dev/null +++ b/plugins/det_post/det_post_node.cpp @@ -0,0 +1,852 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "node.h" + +namespace rk3588 { + +namespace { + +struct ImgSize { + int w = 0; + int h = 0; +}; + +struct Roi01 { + float x = 0.0f; + float y = 0.0f; + float w = 1.0f; + float h = 1.0f; +}; + +static bool ParseRoi01(const SimpleJson& cfg, Roi01& out) { + if (!cfg.IsObject()) return false; + out.x = static_cast(cfg.ValueOr("x", out.x)); + out.y = static_cast(cfg.ValueOr("y", out.y)); + out.w = static_cast(cfg.ValueOr("w", out.w)); + out.h = static_cast(cfg.ValueOr("h", out.h)); + return true; +} + +static bool CenterInRoi(const Detection& det, const Roi01& roi, const ImgSize& img) { + if (img.w <= 0 || img.h <= 0) return true; + const float cx = det.bbox.x + det.bbox.w * 0.5f; + const float cy = det.bbox.y + det.bbox.h * 0.5f; + + const float rx1 = roi.x * img.w; + const float ry1 = roi.y * img.h; + const float rx2 = (roi.x + roi.w) * img.w; + const float ry2 = (roi.y + roi.h) * img.h; + + return (cx >= rx1 && cx <= rx2 && cy >= ry1 && cy <= ry2); +} + +static inline float ClampF(float v, float lo, float hi) { + if (v < lo) return lo; + if (v > hi) return hi; + return v; +} + +static inline int ClampI(int v, int lo, int hi) { + if (v < lo) return lo; + if (v > hi) return hi; + return v; +} + +static inline bool GetRgbAt(const Frame& frame, int x, int y, uint8_t& r, uint8_t& g, uint8_t& b) { + if (!frame.data || frame.width <= 0 || frame.height <= 0) return false; + if (x < 0 || y < 0 || x >= frame.width || y >= frame.height) return false; + + auto plane_ptr = [&](int idx) -> const uint8_t* { + if (idx < 0 || idx >= frame.plane_count) return nullptr; + if (frame.planes[idx].data) return frame.planes[idx].data; + const int off = frame.planes[idx].offset; + if (off < 0) return nullptr; + return frame.data + static_cast(off); + }; + auto plane_stride = [&](int idx, int fallback) -> int { + if (idx >= 0 && idx < frame.plane_count && frame.planes[idx].stride > 0) return frame.planes[idx].stride; + if (frame.stride > 0) return frame.stride; + return fallback; + }; + + if (frame.format == PixelFormat::RGB || frame.format == PixelFormat::BGR) { + const int stride = plane_stride(0, frame.width * 3); + const uint8_t* p = plane_ptr(0); + if (!p) return false; + const size_t idx = static_cast(y) * static_cast(stride) + static_cast(x) * 3; + if (frame.format == PixelFormat::RGB) { + r = p[idx + 0]; + g = p[idx + 1]; + b = p[idx + 2]; + } else { + b = p[idx + 0]; + g = p[idx + 1]; + r = p[idx + 2]; + } + return true; + } + + if (frame.format == PixelFormat::NV12) { + const uint8_t* y0 = plane_ptr(0); + if (!y0) return false; + const int y_stride = plane_stride(0, frame.width); + const uint8_t* uv0 = plane_ptr(1); + const int uv_stride = plane_stride(1, frame.width); + if (!uv0) { + uv0 = frame.data + static_cast(y_stride) * static_cast(frame.height); + } + + const int Y = y0[y * y_stride + x]; + const int uvx = (x / 2) * 2; + const int uvy = (y / 2); + const int U = uv0[uvy * uv_stride + uvx + 0]; + const int V = uv0[uvy * uv_stride + uvx + 1]; + + // BT.601 full range approximation + const float yf = static_cast(Y); + const float uf = static_cast(U) - 128.0f; + const float vf = static_cast(V) - 128.0f; + float rf = yf + 1.402f * vf; + float gf = yf - 0.344136f * uf - 0.714136f * vf; + float bf = yf + 1.772f * uf; + r = static_cast(ClampI(static_cast(std::lround(rf)), 0, 255)); + g = static_cast(ClampI(static_cast(std::lround(gf)), 0, 255)); + b = static_cast(ClampI(static_cast(std::lround(bf)), 0, 255)); + return true; + } + + if (frame.format == PixelFormat::YUV420) { + const uint8_t* y0 = plane_ptr(0); + const uint8_t* u0 = plane_ptr(1); + const uint8_t* v0 = plane_ptr(2); + if (!y0 || !u0 || !v0) return false; + const int y_stride = plane_stride(0, frame.width); + const int u_stride = plane_stride(1, frame.width / 2); + const int v_stride = plane_stride(2, frame.width / 2); + + const int Y = y0[y * y_stride + x]; + const int U = u0[(y / 2) * u_stride + (x / 2)]; + const int V = v0[(y / 2) * v_stride + (x / 2)]; + + const float yf = static_cast(Y); + const float uf = static_cast(U) - 128.0f; + const float vf = static_cast(V) - 128.0f; + float rf = yf + 1.402f * vf; + float gf = yf - 0.344136f * uf - 0.714136f * vf; + float bf = yf + 1.772f * uf; + r = static_cast(ClampI(static_cast(std::lround(rf)), 0, 255)); + g = static_cast(ClampI(static_cast(std::lround(gf)), 0, 255)); + b = static_cast(ClampI(static_cast(std::lround(bf)), 0, 255)); + return true; + } + + return false; +} + +static inline void RgbToHsv(uint8_t r8, uint8_t g8, uint8_t b8, float& h_deg, float& s_255, float& v_255) { + const float r = r8 / 255.0f; + const float g = g8 / 255.0f; + const float b = b8 / 255.0f; + + const float mx = std::max({r, g, b}); + const float mn = std::min({r, g, b}); + const float d = mx - mn; + + float h = 0.0f; + if (d > 1e-6f) { + if (mx == r) { + h = 60.0f * std::fmod(((g - b) / d), 6.0f); + } else if (mx == g) { + h = 60.0f * (((b - r) / d) + 2.0f); + } else { + h = 60.0f * (((r - g) / d) + 4.0f); + } + } + if (h < 0.0f) h += 360.0f; + const float s = (mx <= 1e-6f) ? 0.0f : (d / mx); + h_deg = h; + s_255 = s * 255.0f; + v_255 = mx * 255.0f; +} + +struct Lab { + float L = 0.0f; + float a = 0.0f; + float b = 0.0f; +}; + +static inline float SrgbToLinear(float c01) { + if (c01 <= 0.04045f) return c01 / 12.92f; + return std::pow((c01 + 0.055f) / 1.055f, 2.4f); +} + +static inline float LabF(float t) { + constexpr float d = 6.0f / 29.0f; + constexpr float d3 = d * d * d; + if (t > d3) return std::cbrt(t); + return t / (3.0f * d * d) + 4.0f / 29.0f; +} + +static inline Lab RgbToLab(uint8_t r8, uint8_t g8, uint8_t b8) { + const float r = SrgbToLinear(r8 / 255.0f); + const float g = SrgbToLinear(g8 / 255.0f); + const float b = SrgbToLinear(b8 / 255.0f); + + // sRGB D65 + const float X = 0.4124564f * r + 0.3575761f * g + 0.1804375f * b; + const float Y = 0.2126729f * r + 0.7151522f * g + 0.0721750f * b; + const float Z = 0.0193339f * r + 0.1191920f * g + 0.9503041f * b; + + constexpr float Xn = 0.95047f; + constexpr float Yn = 1.00000f; + constexpr float Zn = 1.08883f; + + const float fx = LabF(X / Xn); + const float fy = LabF(Y / Yn); + const float fz = LabF(Z / Zn); + + Lab out; + out.L = 116.0f * fy - 16.0f; + out.a = 500.0f * (fx - fy); + out.b = 200.0f * (fy - fz); + return out; +} + +struct Crop01 { + float x = 0.0f; + float y = 0.0f; + float w = 1.0f; + float h = 1.0f; +}; + +static bool ParseCrop01(const SimpleJson& cfg, Crop01& out) { + if (!cfg.IsObject()) return false; + out.x = static_cast(cfg.ValueOr("x", out.x)); + out.y = static_cast(cfg.ValueOr("y", out.y)); + out.w = static_cast(cfg.ValueOr("w", out.w)); + out.h = static_cast(cfg.ValueOr("h", out.h)); + out.x = ClampF(out.x, 0.0f, 1.0f); + out.y = ClampF(out.y, 0.0f, 1.0f); + out.w = ClampF(out.w, 0.0f, 1.0f); + out.h = ClampF(out.h, 0.0f, 1.0f); + return true; +} + +class IProcessor { +public: + virtual ~IProcessor() = default; + virtual const char* Type() const = 0; + virtual bool Apply(const Frame& frame, const ImgSize& img, Detection& det) const = 0; +}; + +class ConfGate final : public IProcessor { +public: + explicit ConfGate(float conf_min) : conf_min_(conf_min) {} + const char* Type() const override { return "conf_gate"; } + bool Apply(const Frame&, const ImgSize&, Detection& det) const override { + return det.score >= conf_min_; + } + +private: + float conf_min_ = 0.0f; +}; + +class BboxGate final : public IProcessor { +public: + BboxGate(float min_w, float min_h) : min_w_(min_w), min_h_(min_h) {} + const char* Type() const override { return "bbox_gate"; } + bool Apply(const Frame&, const ImgSize&, Detection& det) const override { + return det.bbox.w >= min_w_ && det.bbox.h >= min_h_; + } + +private: + float min_w_ = 0.0f; + float min_h_ = 0.0f; +}; + +class RoiGate final : public IProcessor { +public: + RoiGate(Roi01 roi, bool center_mode) : roi_(roi), center_mode_(center_mode) {} + const char* Type() const override { return "roi_gate"; } + bool Apply(const Frame&, const ImgSize& img, Detection& det) const override { + if (center_mode_) return CenterInRoi(det, roi_, img); + return CenterInRoi(det, roi_, img); + } + +private: + Roi01 roi_{}; + bool center_mode_ = true; +}; + +class HsvRatio final : public IProcessor { +public: + struct Range { + float h_min = 0.0f; + float h_max = 360.0f; + float s_min = 0.0f; + float s_max = 255.0f; + float v_min = 0.0f; + float v_max = 255.0f; + }; + + HsvRatio(Range range, float ratio_min, int max_samples, int step, int min_samples, Crop01 crop) + : range_(range), ratio_min_(ratio_min), max_samples_(max_samples), step_(step), + min_samples_(min_samples), crop_(crop) {} + + const char* Type() const override { return "hsv_ratio"; } + + bool Apply(const Frame& frame, const ImgSize&, Detection& det) const override { + if (!frame.data || frame.width <= 0 || frame.height <= 0) return true; + + int x0 = static_cast(std::floor(det.bbox.x)); + int y0 = static_cast(std::floor(det.bbox.y)); + int w = static_cast(std::floor(det.bbox.w)); + int h = static_cast(std::floor(det.bbox.h)); + if (w <= 1 || h <= 1) return true; + + x0 = ClampI(x0, 0, frame.width - 1); + y0 = ClampI(y0, 0, frame.height - 1); + w = ClampI(w, 1, frame.width - x0); + h = ClampI(h, 1, frame.height - y0); + + // Apply crop within bbox + int cx0 = x0 + static_cast(std::floor(crop_.x * w)); + int cy0 = y0 + static_cast(std::floor(crop_.y * h)); + int cw = static_cast(std::floor(crop_.w * w)); + int ch = static_cast(std::floor(crop_.h * h)); + if (cw <= 1 || ch <= 1) return true; + cx0 = ClampI(cx0, 0, frame.width - 1); + cy0 = ClampI(cy0, 0, frame.height - 1); + cw = ClampI(cw, 1, frame.width - cx0); + ch = ClampI(ch, 1, frame.height - cy0); + + int step = step_; + if (step <= 0) { + const double area = static_cast(cw) * static_cast(ch); + const double target = std::max(1, max_samples_); + step = static_cast(std::floor(std::sqrt(area / target))); + if (step < 1) step = 1; + } + + int total = 0; + int matched = 0; + for (int yy = cy0; yy < cy0 + ch; yy += step) { + for (int xx = cx0; xx < cx0 + cw; xx += step) { + uint8_t r, g, b; + if (!GetRgbAt(frame, xx, yy, r, g, b)) continue; + float hdeg, s255, v255; + RgbToHsv(r, g, b, hdeg, s255, v255); + ++total; + if (InRange(hdeg, s255, v255)) { + ++matched; + } + if (max_samples_ > 0 && total >= max_samples_) break; + } + if (max_samples_ > 0 && total >= max_samples_) break; + } + + if (total < std::max(1, min_samples_)) return true; + const float ratio = static_cast(matched) / static_cast(total); + return ratio >= ratio_min_; + } + +private: + bool InRange(float hdeg, float s255, float v255) const { + const bool h_ok = [&] { + float hmin = range_.h_min; + float hmax = range_.h_max; + // Normalize + while (hmin < 0.0f) hmin += 360.0f; + while (hmax < 0.0f) hmax += 360.0f; + while (hmin >= 360.0f) hmin -= 360.0f; + while (hmax >= 360.0f) hmax -= 360.0f; + if (hmin <= hmax) { + return (hdeg >= hmin && hdeg <= hmax); + } + // wrap-around range + return (hdeg >= hmin || hdeg <= hmax); + }(); + + return h_ok && (s255 >= range_.s_min && s255 <= range_.s_max) && (v255 >= range_.v_min && v255 <= range_.v_max); + } + + Range range_{}; + float ratio_min_ = 0.0f; + int max_samples_ = 2048; + int step_ = 0; + int min_samples_ = 32; + Crop01 crop_{}; +}; + +class LabKmeans final : public IProcessor { +public: + LabKmeans(int k, int iters, int max_samples, int step, int min_samples, float l_max, float ratio_min, Crop01 crop) + : k_(k), iters_(iters), max_samples_(max_samples), step_(step), min_samples_(min_samples), + l_max_(l_max), ratio_min_(ratio_min), crop_(crop) {} + + const char* Type() const override { return "lab_kmeans"; } + + bool Apply(const Frame& frame, const ImgSize&, Detection& det) const override { + if (!frame.data || frame.width <= 0 || frame.height <= 0) return true; + + int x0 = static_cast(std::floor(det.bbox.x)); + int y0 = static_cast(std::floor(det.bbox.y)); + int w = static_cast(std::floor(det.bbox.w)); + int h = static_cast(std::floor(det.bbox.h)); + if (w <= 1 || h <= 1) return true; + + x0 = ClampI(x0, 0, frame.width - 1); + y0 = ClampI(y0, 0, frame.height - 1); + w = ClampI(w, 1, frame.width - x0); + h = ClampI(h, 1, frame.height - y0); + + int cx0 = x0 + static_cast(std::floor(crop_.x * w)); + int cy0 = y0 + static_cast(std::floor(crop_.y * h)); + int cw = static_cast(std::floor(crop_.w * w)); + int ch = static_cast(std::floor(crop_.h * h)); + if (cw <= 1 || ch <= 1) return true; + cx0 = ClampI(cx0, 0, frame.width - 1); + cy0 = ClampI(cy0, 0, frame.height - 1); + cw = ClampI(cw, 1, frame.width - cx0); + ch = ClampI(ch, 1, frame.height - cy0); + + int step = step_; + if (step <= 0) { + const double area = static_cast(cw) * static_cast(ch); + const double target = std::max(1, max_samples_); + step = static_cast(std::floor(std::sqrt(area / target))); + if (step < 1) step = 1; + } + + std::vector pts; + pts.reserve(static_cast(std::max(0, max_samples_))); + + for (int yy = cy0; yy < cy0 + ch; yy += step) { + for (int xx = cx0; xx < cx0 + cw; xx += step) { + uint8_t r, g, b; + if (!GetRgbAt(frame, xx, yy, r, g, b)) continue; + pts.push_back(RgbToLab(r, g, b)); + if (max_samples_ > 0 && static_cast(pts.size()) >= max_samples_) break; + } + if (max_samples_ > 0 && static_cast(pts.size()) >= max_samples_) break; + } + + if (static_cast(pts.size()) < std::max(1, min_samples_)) return true; + + const int k = ClampI(k_, 2, 8); + std::vector centroids(static_cast(k)); + std::vector assign(pts.size(), 0); + std::vector counts(static_cast(k), 0); + + // deterministic seed based on bbox and class + uint64_t seed = (static_cast(det.cls_id) << 32) ^ + (static_cast(x0) << 16) ^ (static_cast(y0) << 1) ^ + (static_cast(w) << 8) ^ static_cast(h); + std::mt19937 rng(static_cast(seed ^ (seed >> 32))); + std::uniform_int_distribution uni(0, pts.size() - 1); + + // init: random distinct points (best-effort) + for (int i = 0; i < k; ++i) { + centroids[static_cast(i)] = pts[uni(rng)]; + } + + for (int iter = 0; iter < std::max(1, iters_); ++iter) { + std::fill(counts.begin(), counts.end(), 0); + std::vector sums(static_cast(k)); + + for (size_t i = 0; i < pts.size(); ++i) { + const Lab& p = pts[i]; + int best = 0; + float best_d = Dist2(p, centroids[0]); + for (int c = 1; c < k; ++c) { + const float d = Dist2(p, centroids[static_cast(c)]); + if (d < best_d) { + best_d = d; + best = c; + } + } + assign[i] = best; + counts[static_cast(best)] += 1; + sums[static_cast(best)].L += p.L; + sums[static_cast(best)].a += p.a; + sums[static_cast(best)].b += p.b; + } + + for (int c = 0; c < k; ++c) { + const int cnt = counts[static_cast(c)]; + if (cnt <= 0) { + centroids[static_cast(c)] = pts[uni(rng)]; + continue; + } + Lab m; + m.L = sums[static_cast(c)].L / cnt; + m.a = sums[static_cast(c)].a / cnt; + m.b = sums[static_cast(c)].b / cnt; + centroids[static_cast(c)] = m; + } + } + + // pick darkest centroid + int darkest = 0; + for (int c = 1; c < k; ++c) { + if (centroids[static_cast(c)].L < centroids[static_cast(darkest)].L) darkest = c; + } + const float darkest_L = centroids[static_cast(darkest)].L; + if (darkest_L > l_max_) return false; + + const float ratio = static_cast(counts[static_cast(darkest)]) / static_cast(pts.size()); + return ratio >= ratio_min_; + } + +private: + static float Dist2(const Lab& p, const Lab& c) { + const float dL = p.L - c.L; + const float da = p.a - c.a; + const float db = p.b - c.b; + return dL * dL + da * da + db * db; + } + + int k_ = 3; + int iters_ = 8; + int max_samples_ = 2048; + int step_ = 0; + int min_samples_ = 64; + float l_max_ = 50.0f; + float ratio_min_ = 0.4f; + Crop01 crop_{}; +}; + +using FactoryFn = std::function(const SimpleJson& cfg, std::string& err)>; + +static const std::unordered_map& Registry() { + static const std::unordered_map kReg = { + {"conf_gate", + [](const SimpleJson& cfg, std::string&) -> std::unique_ptr { + const float conf_min = static_cast(cfg.ValueOr("conf_min", 0.0)); + return std::make_unique(conf_min); + }}, + {"bbox_gate", + [](const SimpleJson& cfg, std::string&) -> std::unique_ptr { + const float min_w = static_cast(cfg.ValueOr("min_w", 0.0)); + const float min_h = static_cast(cfg.ValueOr("min_h", 0.0)); + return std::make_unique(min_w, min_h); + }}, + {"roi_gate", + [](const SimpleJson& cfg, std::string& err) -> std::unique_ptr { + Roi01 roi; + if (const SimpleJson* r = cfg.Find("roi")) { + if (!ParseRoi01(*r, roi)) { + err = "roi_gate.roi must be object"; + return nullptr; + } + } + const std::string mode = cfg.ValueOr("mode", "center"); + const bool center_mode = (mode == "center"); + return std::make_unique(roi, center_mode); + }}, + {"hsv_ratio", + [](const SimpleJson& cfg, std::string&) -> std::unique_ptr { + HsvRatio::Range r; + r.h_min = static_cast(cfg.ValueOr("h_min", r.h_min)); + r.h_max = static_cast(cfg.ValueOr("h_max", r.h_max)); + r.s_min = static_cast(cfg.ValueOr("s_min", r.s_min)); + r.s_max = static_cast(cfg.ValueOr("s_max", r.s_max)); + r.v_min = static_cast(cfg.ValueOr("v_min", r.v_min)); + r.v_max = static_cast(cfg.ValueOr("v_max", r.v_max)); + + const float ratio_min = static_cast(cfg.ValueOr("ratio_min", 0.0)); + const int max_samples = cfg.ValueOr("max_samples", 2048); + const int step = cfg.ValueOr("step", 0); + const int min_samples = cfg.ValueOr("min_samples", 32); + Crop01 crop; + if (const SimpleJson* c = cfg.Find("crop")) (void)ParseCrop01(*c, crop); + return std::make_unique(r, ratio_min, max_samples, step, min_samples, crop); + }}, + {"lab_kmeans", + [](const SimpleJson& cfg, std::string&) -> std::unique_ptr { + const int k = cfg.ValueOr("k", 3); + const int iters = cfg.ValueOr("iters", 8); + const int max_samples = cfg.ValueOr("max_samples", 2048); + const int step = cfg.ValueOr("step", 0); + const int min_samples = cfg.ValueOr("min_samples", 64); + const float l_max = static_cast(cfg.ValueOr("l_max", 50.0)); + const float ratio_min = static_cast(cfg.ValueOr("ratio_min", 0.4)); + Crop01 crop; + if (const SimpleJson* c = cfg.Find("crop")) (void)ParseCrop01(*c, crop); + return std::make_unique(k, iters, max_samples, step, min_samples, l_max, ratio_min, crop); + }}, + }; + return kReg; +} + +struct ClassPipeline { + bool enable = false; + std::vector> processors; +}; + +struct ConfigSnapshot { + std::string id; + ClassPipeline def; + std::unordered_map per_class; +}; + +static bool BuildPipeline(const SimpleJson& cfg, ClassPipeline& out, std::string& err) { + out.enable = cfg.ValueOr("enable", out.enable); + out.processors.clear(); + + const SimpleJson* ps = cfg.Find("processors"); + if (!ps) return true; + if (!ps->IsArray()) { + err = "processors must be array"; + return false; + } + + for (const auto& p : ps->AsArray()) { + if (!p.IsObject()) { + err = "processor entry must be object"; + return false; + } + const std::string type = p.ValueOr("type", ""); + if (type.empty()) { + err = "processor missing 'type'"; + return false; + } + auto it = Registry().find(type); + if (it == Registry().end()) { + err = "unknown processor type: " + type; + return false; + } + std::string perr; + auto proc = it->second(p, perr); + if (!proc) { + err = perr.empty() ? ("failed to create processor: " + type) : perr; + return false; + } + out.processors.push_back(std::move(proc)); + } + return true; +} + +static bool BuildConfigSnapshot(const SimpleJson& config, std::shared_ptr& out, + std::string& err) { + auto snap = std::make_shared(); + snap->id = config.ValueOr("id", "det_post"); + + // default pipeline + if (const SimpleJson* d = config.Find("default")) { + if (!d->IsObject()) { + err = "default must be object"; + return false; + } + snap->def.enable = d->ValueOr("enable", false); + if (!BuildPipeline(*d, snap->def, err)) return false; + } else { + snap->def.enable = false; + } + + if (const SimpleJson* pc = config.Find("per_class")) { + if (!pc->IsArray()) { + err = "per_class must be array"; + return false; + } + for (const auto& it : pc->AsArray()) { + if (!it.IsObject()) { + err = "per_class entry must be object"; + return false; + } + const int cls = it.ValueOr("class_id", -1); + if (cls < 0) { + err = "per_class.class_id must be >= 0"; + return false; + } + ClassPipeline cp; + cp.enable = it.ValueOr("enable", false); + if (!BuildPipeline(it, cp, err)) return false; + snap->per_class.emplace(cls, std::move(cp)); + } + } + + out = std::move(snap); + return true; +} + +} // namespace + +class DetPostNode final : public INode { +public: + std::string Id() const override { return id_; } + std::string Type() const override { return "det_post"; } + + bool Init(const SimpleJson& config, const NodeContext& ctx) override { + std::shared_ptr snap; + std::string err; + if (!BuildConfigSnapshot(config, snap, err)) { + std::cerr << "[det_post] invalid config: " << err << "\n"; + return false; + } + id_ = snap->id; + + input_queue_ = ctx.input_queue; + if (!input_queue_) { + std::cerr << "[det_post] no input queue for node " << id_ << "\n"; + return false; + } + if (ctx.output_queues.empty()) { + std::cerr << "[det_post] no output queue for node " << id_ << "\n"; + return false; + } + output_queues_ = ctx.output_queues; + + { + std::lock_guard lock(mu_); + cfg_ = std::move(snap); + } + return true; + } + + bool Start() override { + std::cout << "[det_post] started\n"; + return true; + } + + void Stop() override { + std::cout << "[det_post] stopped\n"; + } + + bool UpdateConfig(const SimpleJson& new_config) override { + std::shared_ptr snap; + std::string err; + if (!BuildConfigSnapshot(new_config, snap, err)) { + std::cerr << "[det_post] UpdateConfig rejected: " << err << "\n"; + return false; + } + if (!id_.empty() && !snap->id.empty() && snap->id != id_) { + return false; + } + { + std::lock_guard lock(mu_); + cfg_ = std::move(snap); + } + return true; + } + + bool GetCustomMetrics(SimpleJson& out) const override { + SimpleJson::Object o; + o["processed"] = SimpleJson(static_cast(processed_.load())); + o["dropped"] = SimpleJson(static_cast(dropped_.load())); + { + std::lock_guard lock(mu_); + if (cfg_) { + o["per_class_count"] = SimpleJson(static_cast(cfg_->per_class.size())); + o["default_enable"] = SimpleJson(cfg_->def.enable); + } + } + out = SimpleJson(std::move(o)); + return true; + } + + NodeStatus Process(FramePtr frame) override { + if (!frame) return NodeStatus::DROP; + + std::shared_ptr cfg; + { + std::lock_guard lock(mu_); + cfg = cfg_; + } + if (!cfg) { + PushToDownstream(frame); + processed_.fetch_add(1); + return NodeStatus::OK; + } + + if (!frame->det || frame->det->items.empty()) { + PushToDownstream(frame); + processed_.fetch_add(1); + return NodeStatus::OK; + } + + const auto& in = frame->det; + ImgSize img; + img.w = (in->img_w > 0) ? in->img_w : frame->width; + img.h = (in->img_h > 0) ? in->img_h : frame->height; + + auto out_det = std::make_shared(); + out_det->img_w = in->img_w; + out_det->img_h = in->img_h; + out_det->model_name = in->model_name; + out_det->items.reserve(in->items.size()); + + uint64_t dropped_local = 0; + + for (const auto& det0 : in->items) { + Detection det = det0; + + const ClassPipeline* pipeline = nullptr; + auto it = cfg->per_class.find(det.cls_id); + if (it != cfg->per_class.end()) { + pipeline = &it->second; + } else if (cfg->def.enable) { + pipeline = &cfg->def; + } + + bool keep = true; + if (pipeline && pipeline->enable) { + for (const auto& p : pipeline->processors) { + if (!p) continue; + if (!p->Apply(*frame, img, det)) { + keep = false; + break; + } + } + } + + if (keep) { + out_det->items.push_back(det); + } else { + ++dropped_local; + } + } + + frame->det = std::move(out_det); + + if (dropped_local) dropped_.fetch_add(dropped_local); + processed_.fetch_add(1); + PushToDownstream(frame); + return NodeStatus::OK; + } + +private: + void PushToDownstream(FramePtr frame) { + for (auto& q : output_queues_) { + q->Push(frame); + } + } + + std::string id_; + + std::shared_ptr> input_queue_; + std::vector>> output_queues_; + + mutable std::mutex mu_; + std::shared_ptr cfg_; + + std::atomic processed_{0}; + std::atomic dropped_{0}; +}; + +REGISTER_NODE(DetPostNode, "det_post"); + +} // namespace rk3588