860 lines
30 KiB
C++
860 lines
30 KiB
C++
#include <atomic>
|
|
#include <atomic>
|
|
#include <algorithm>
|
|
#include <cstdint>
|
|
#include <cmath>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <random>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "node.h"
|
|
#include "utils/logger.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<float>(cfg.ValueOr<double>("x", out.x));
|
|
out.y = static_cast<float>(cfg.ValueOr<double>("y", out.y));
|
|
out.w = static_cast<float>(cfg.ValueOr<double>("w", out.w));
|
|
out.h = static_cast<float>(cfg.ValueOr<double>("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<size_t>(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<size_t>(y) * static_cast<size_t>(stride) + static_cast<size_t>(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<size_t>(y_stride) * static_cast<size_t>(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<float>(Y);
|
|
const float uf = static_cast<float>(U) - 128.0f;
|
|
const float vf = static_cast<float>(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<uint8_t>(ClampI(static_cast<int>(std::lround(rf)), 0, 255));
|
|
g = static_cast<uint8_t>(ClampI(static_cast<int>(std::lround(gf)), 0, 255));
|
|
b = static_cast<uint8_t>(ClampI(static_cast<int>(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<float>(Y);
|
|
const float uf = static_cast<float>(U) - 128.0f;
|
|
const float vf = static_cast<float>(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<uint8_t>(ClampI(static_cast<int>(std::lround(rf)), 0, 255));
|
|
g = static_cast<uint8_t>(ClampI(static_cast<int>(std::lround(gf)), 0, 255));
|
|
b = static_cast<uint8_t>(ClampI(static_cast<int>(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<float>(cfg.ValueOr<double>("x", out.x));
|
|
out.y = static_cast<float>(cfg.ValueOr<double>("y", out.y));
|
|
out.w = static_cast<float>(cfg.ValueOr<double>("w", out.w));
|
|
out.h = static_cast<float>(cfg.ValueOr<double>("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<int>(std::floor(det.bbox.x));
|
|
int y0 = static_cast<int>(std::floor(det.bbox.y));
|
|
int w = static_cast<int>(std::floor(det.bbox.w));
|
|
int h = static_cast<int>(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<int>(std::floor(crop_.x * w));
|
|
int cy0 = y0 + static_cast<int>(std::floor(crop_.y * h));
|
|
int cw = static_cast<int>(std::floor(crop_.w * w));
|
|
int ch = static_cast<int>(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<double>(cw) * static_cast<double>(ch);
|
|
const double target = std::max(1, max_samples_);
|
|
step = static_cast<int>(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<float>(matched) / static_cast<float>(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<int>(std::floor(det.bbox.x));
|
|
int y0 = static_cast<int>(std::floor(det.bbox.y));
|
|
int w = static_cast<int>(std::floor(det.bbox.w));
|
|
int h = static_cast<int>(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<int>(std::floor(crop_.x * w));
|
|
int cy0 = y0 + static_cast<int>(std::floor(crop_.y * h));
|
|
int cw = static_cast<int>(std::floor(crop_.w * w));
|
|
int ch = static_cast<int>(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<double>(cw) * static_cast<double>(ch);
|
|
const double target = std::max(1, max_samples_);
|
|
step = static_cast<int>(std::floor(std::sqrt(area / target)));
|
|
if (step < 1) step = 1;
|
|
}
|
|
|
|
std::vector<Lab> pts;
|
|
pts.reserve(static_cast<size_t>(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<int>(pts.size()) >= max_samples_) break;
|
|
}
|
|
if (max_samples_ > 0 && static_cast<int>(pts.size()) >= max_samples_) break;
|
|
}
|
|
|
|
if (static_cast<int>(pts.size()) < std::max(1, min_samples_)) return true;
|
|
|
|
const int k = ClampI(k_, 2, 8);
|
|
std::vector<Lab> centroids(static_cast<size_t>(k));
|
|
std::vector<int> assign(pts.size(), 0);
|
|
std::vector<int> counts(static_cast<size_t>(k), 0);
|
|
|
|
// deterministic seed based on bbox and class
|
|
uint64_t seed = (static_cast<uint64_t>(det.cls_id) << 32) ^
|
|
(static_cast<uint64_t>(x0) << 16) ^ (static_cast<uint64_t>(y0) << 1) ^
|
|
(static_cast<uint64_t>(w) << 8) ^ static_cast<uint64_t>(h);
|
|
std::mt19937 rng(static_cast<uint32_t>(seed ^ (seed >> 32)));
|
|
std::uniform_int_distribution<size_t> uni(0, pts.size() - 1);
|
|
|
|
// init: random distinct points (best-effort)
|
|
for (int i = 0; i < k; ++i) {
|
|
centroids[static_cast<size_t>(i)] = pts[uni(rng)];
|
|
}
|
|
|
|
for (int iter = 0; iter < std::max(1, iters_); ++iter) {
|
|
std::fill(counts.begin(), counts.end(), 0);
|
|
std::vector<Lab> sums(static_cast<size_t>(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<size_t>(c)]);
|
|
if (d < best_d) {
|
|
best_d = d;
|
|
best = c;
|
|
}
|
|
}
|
|
assign[i] = best;
|
|
counts[static_cast<size_t>(best)] += 1;
|
|
sums[static_cast<size_t>(best)].L += p.L;
|
|
sums[static_cast<size_t>(best)].a += p.a;
|
|
sums[static_cast<size_t>(best)].b += p.b;
|
|
}
|
|
|
|
for (int c = 0; c < k; ++c) {
|
|
const int cnt = counts[static_cast<size_t>(c)];
|
|
if (cnt <= 0) {
|
|
centroids[static_cast<size_t>(c)] = pts[uni(rng)];
|
|
continue;
|
|
}
|
|
Lab m;
|
|
m.L = sums[static_cast<size_t>(c)].L / cnt;
|
|
m.a = sums[static_cast<size_t>(c)].a / cnt;
|
|
m.b = sums[static_cast<size_t>(c)].b / cnt;
|
|
centroids[static_cast<size_t>(c)] = m;
|
|
}
|
|
}
|
|
|
|
// pick darkest centroid
|
|
int darkest = 0;
|
|
for (int c = 1; c < k; ++c) {
|
|
if (centroids[static_cast<size_t>(c)].L < centroids[static_cast<size_t>(darkest)].L) darkest = c;
|
|
}
|
|
const float darkest_L = centroids[static_cast<size_t>(darkest)].L;
|
|
if (darkest_L > l_max_) return false;
|
|
|
|
const float ratio = static_cast<float>(counts[static_cast<size_t>(darkest)]) / static_cast<float>(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<std::unique_ptr<IProcessor>(const SimpleJson& cfg, std::string& err)>;
|
|
|
|
static const std::unordered_map<std::string, FactoryFn>& Registry() {
|
|
static const std::unordered_map<std::string, FactoryFn> kReg = {
|
|
{"conf_gate",
|
|
[](const SimpleJson& cfg, std::string&) -> std::unique_ptr<IProcessor> {
|
|
const float conf_min = static_cast<float>(cfg.ValueOr<double>("conf_min", 0.0));
|
|
return std::make_unique<ConfGate>(conf_min);
|
|
}},
|
|
{"bbox_gate",
|
|
[](const SimpleJson& cfg, std::string&) -> std::unique_ptr<IProcessor> {
|
|
const float min_w = static_cast<float>(cfg.ValueOr<double>("min_w", 0.0));
|
|
const float min_h = static_cast<float>(cfg.ValueOr<double>("min_h", 0.0));
|
|
return std::make_unique<BboxGate>(min_w, min_h);
|
|
}},
|
|
{"roi_gate",
|
|
[](const SimpleJson& cfg, std::string& err) -> std::unique_ptr<IProcessor> {
|
|
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<std::string>("mode", "center");
|
|
const bool center_mode = (mode == "center");
|
|
return std::make_unique<RoiGate>(roi, center_mode);
|
|
}},
|
|
{"hsv_ratio",
|
|
[](const SimpleJson& cfg, std::string&) -> std::unique_ptr<IProcessor> {
|
|
HsvRatio::Range r;
|
|
r.h_min = static_cast<float>(cfg.ValueOr<double>("h_min", r.h_min));
|
|
r.h_max = static_cast<float>(cfg.ValueOr<double>("h_max", r.h_max));
|
|
r.s_min = static_cast<float>(cfg.ValueOr<double>("s_min", r.s_min));
|
|
r.s_max = static_cast<float>(cfg.ValueOr<double>("s_max", r.s_max));
|
|
r.v_min = static_cast<float>(cfg.ValueOr<double>("v_min", r.v_min));
|
|
r.v_max = static_cast<float>(cfg.ValueOr<double>("v_max", r.v_max));
|
|
|
|
const float ratio_min = static_cast<float>(cfg.ValueOr<double>("ratio_min", 0.0));
|
|
const int max_samples = cfg.ValueOr<int>("max_samples", 2048);
|
|
const int step = cfg.ValueOr<int>("step", 0);
|
|
const int min_samples = cfg.ValueOr<int>("min_samples", 32);
|
|
Crop01 crop;
|
|
if (const SimpleJson* c = cfg.Find("crop")) (void)ParseCrop01(*c, crop);
|
|
return std::make_unique<HsvRatio>(r, ratio_min, max_samples, step, min_samples, crop);
|
|
}},
|
|
{"lab_kmeans",
|
|
[](const SimpleJson& cfg, std::string&) -> std::unique_ptr<IProcessor> {
|
|
const int k = cfg.ValueOr<int>("k", 3);
|
|
const int iters = cfg.ValueOr<int>("iters", 8);
|
|
const int max_samples = cfg.ValueOr<int>("max_samples", 2048);
|
|
const int step = cfg.ValueOr<int>("step", 0);
|
|
const int min_samples = cfg.ValueOr<int>("min_samples", 64);
|
|
const float l_max = static_cast<float>(cfg.ValueOr<double>("l_max", 50.0));
|
|
const float ratio_min = static_cast<float>(cfg.ValueOr<double>("ratio_min", 0.4));
|
|
Crop01 crop;
|
|
if (const SimpleJson* c = cfg.Find("crop")) (void)ParseCrop01(*c, crop);
|
|
return std::make_unique<LabKmeans>(k, iters, max_samples, step, min_samples, l_max, ratio_min, crop);
|
|
}},
|
|
};
|
|
return kReg;
|
|
}
|
|
|
|
struct ClassPipeline {
|
|
bool enable = false;
|
|
std::vector<std::unique_ptr<IProcessor>> processors;
|
|
};
|
|
|
|
struct ConfigSnapshot {
|
|
std::string id;
|
|
ClassPipeline def;
|
|
std::unordered_map<int, ClassPipeline> per_class;
|
|
};
|
|
|
|
static bool BuildPipeline(const SimpleJson& cfg, ClassPipeline& out, std::string& err) {
|
|
out.enable = cfg.ValueOr<bool>("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<std::string>("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<const ConfigSnapshot>& out,
|
|
std::string& err) {
|
|
auto snap = std::make_shared<ConfigSnapshot>();
|
|
snap->id = config.ValueOr<std::string>("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<bool>("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<int>("class_id", -1);
|
|
if (cls < 0) {
|
|
err = "per_class.class_id must be >= 0";
|
|
return false;
|
|
}
|
|
ClassPipeline cp;
|
|
cp.enable = it.ValueOr<bool>("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<const ConfigSnapshot> snap;
|
|
std::string err;
|
|
if (!BuildConfigSnapshot(config, snap, err)) {
|
|
LogError("[det_post] invalid config: " + err);
|
|
return false;
|
|
}
|
|
id_ = snap->id;
|
|
|
|
input_queue_ = ctx.input_queue;
|
|
if (!input_queue_) {
|
|
LogError("[det_post] no input queue for node " + id_);
|
|
return false;
|
|
}
|
|
if (ctx.output_queues.empty()) {
|
|
LogError("[det_post] no output queue for node " + id_);
|
|
return false;
|
|
}
|
|
output_queues_ = ctx.output_queues;
|
|
|
|
{
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
cfg_ = std::move(snap);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool Start() override {
|
|
LogInfo("[det_post] started");
|
|
return true;
|
|
}
|
|
|
|
void Stop() override {
|
|
LogInfo("[det_post] stopped");
|
|
}
|
|
|
|
bool UpdateConfig(const SimpleJson& new_config) override {
|
|
std::shared_ptr<const ConfigSnapshot> snap;
|
|
std::string err;
|
|
if (!BuildConfigSnapshot(new_config, snap, err)) {
|
|
LogWarn("[det_post] UpdateConfig rejected: " + err);
|
|
return false;
|
|
}
|
|
if (!id_.empty() && !snap->id.empty() && snap->id != id_) {
|
|
return false;
|
|
}
|
|
{
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
cfg_ = std::move(snap);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool GetCustomMetrics(SimpleJson& out) const override {
|
|
SimpleJson::Object o;
|
|
o["processed"] = SimpleJson(static_cast<double>(processed_.load()));
|
|
o["dropped"] = SimpleJson(static_cast<double>(dropped_.load()));
|
|
{
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
if (cfg_) {
|
|
o["per_class_count"] = SimpleJson(static_cast<double>(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<const ConfigSnapshot> cfg;
|
|
{
|
|
std::lock_guard<std::mutex> lock(mu_);
|
|
cfg = cfg_;
|
|
}
|
|
if (!cfg) {
|
|
PushToDownstream(frame);
|
|
processed_.fetch_add(1);
|
|
return NodeStatus::OK;
|
|
}
|
|
|
|
// Fast path: no enabled pipelines configured.
|
|
if (!cfg->def.enable && cfg->per_class.empty()) {
|
|
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<DetectionResult>();
|
|
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<SpscQueue<FramePtr>> input_queue_;
|
|
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
|
|
|
|
mutable std::mutex mu_;
|
|
std::shared_ptr<const ConfigSnapshot> cfg_;
|
|
|
|
std::atomic<uint64_t> processed_{0};
|
|
std::atomic<uint64_t> dropped_{0};
|
|
};
|
|
|
|
REGISTER_NODE(DetPostNode, "det_post");
|
|
|
|
} // namespace rk3588
|