463 lines
16 KiB
C++
463 lines
16 KiB
C++
/**
|
|
* ai_shoe_det - 鞋子检测节点（支持滑动窗口）
|
|
*
|
|
* 基于 ai_yolo 实现，添加多窗口检测支持
|
|
*/
|
|
|
|
#include <algorithm>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "ai_scheduler.h"
|
|
#include "frame/frame.h"
|
|
#include "hw/i_infer_backend.h"
|
|
#include "node.h"
|
|
#include "utils/logger.h"
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
#include "rknn_api.h"
|
|
#endif
|
|
|
|
namespace rk3588 {
|
|
|
|
struct DetWindow {
|
|
int x, y, w, h;
|
|
};
|
|
|
|
struct DetBox {
|
|
float x, y, w, h;
|
|
float conf;
|
|
int class_id;
|
|
};
|
|
|
|
struct DynamicRoiConfig {
|
|
bool enable = false;
|
|
int person_class_id = 0;
|
|
int shoe_class_id = 0;
|
|
int max_rois = 8;
|
|
int min_person_height = 0;
|
|
float x_offset = -0.15f;
|
|
float y_offset = 0.72f;
|
|
float width_scale = 1.30f;
|
|
float height_scale = 0.38f;
|
|
};
|
|
|
|
class AiShoeDetNode : public INode {
|
|
public:
|
|
std::string Id() const override { return id_; }
|
|
std::string Type() const override { return "ai_shoe_det"; }
|
|
|
|
bool Init(const SimpleJson& config, const NodeContext& ctx) override {
|
|
id_ = config.ValueOr<std::string>("id", "shoe_det");
|
|
model_path_ = config.ValueOr<std::string>("model_path",
|
|
"./models/shoe_detector.rknn");
|
|
|
|
model_w_ = config.ValueOr<int>("model_w", 640);
|
|
model_h_ = config.ValueOr<int>("model_h", 640);
|
|
conf_thresh_ = config.ValueOr<float>("conf", 0.25f);
|
|
nms_thresh_ = config.ValueOr<float>("nms", 0.45f);
|
|
append_detections_ = config.ValueOr<bool>("append_detections", false);
|
|
|
|
infer_interval_ms_ = std::max<int64_t>(
|
|
0, static_cast<int64_t>(config.ValueOr<int>("infer_interval_ms", 0)));
|
|
if (infer_interval_ms_ <= 0) {
|
|
const double infer_fps = config.ValueOr<double>("infer_fps", 0.0);
|
|
if (infer_fps > 0.0) {
|
|
infer_interval_ms_ = static_cast<int64_t>(1000.0 / infer_fps);
|
|
if (infer_interval_ms_ < 1) infer_interval_ms_ = 1;
|
|
}
|
|
}
|
|
|
|
if (const SimpleJson* dyn = config.Find("dynamic_roi"); dyn && dyn->IsObject()) {
|
|
dynamic_roi_.enable = dyn->ValueOr<bool>("enable", false);
|
|
dynamic_roi_.person_class_id =
|
|
dyn->ValueOr<int>("person_class_id", dynamic_roi_.person_class_id);
|
|
dynamic_roi_.shoe_class_id =
|
|
dyn->ValueOr<int>("shoe_class_id", dynamic_roi_.shoe_class_id);
|
|
dynamic_roi_.max_rois = std::max(1, dyn->ValueOr<int>("max_rois", dynamic_roi_.max_rois));
|
|
dynamic_roi_.min_person_height =
|
|
std::max(0, dyn->ValueOr<int>("min_person_height", dynamic_roi_.min_person_height));
|
|
dynamic_roi_.x_offset = dyn->ValueOr<float>("x_offset", dynamic_roi_.x_offset);
|
|
dynamic_roi_.y_offset = dyn->ValueOr<float>("y_offset", dynamic_roi_.y_offset);
|
|
dynamic_roi_.width_scale = dyn->ValueOr<float>("width_scale", dynamic_roi_.width_scale);
|
|
dynamic_roi_.height_scale = dyn->ValueOr<float>("height_scale", dynamic_roi_.height_scale);
|
|
}
|
|
|
|
// 解析窗口配置
|
|
windows_.clear();
|
|
if (const SimpleJson* win_arr = config.Find("windows"); win_arr && win_arr->IsArray()) {
|
|
for (const auto& w : win_arr->AsArray()) {
|
|
if (w.IsObject()) {
|
|
DetWindow win;
|
|
win.x = w.ValueOr<int>("x", 0);
|
|
win.y = w.ValueOr<int>("y", 0);
|
|
win.w = w.ValueOr<int>("w", 640);
|
|
win.h = w.ValueOr<int>("h", 640);
|
|
windows_.push_back(win);
|
|
}
|
|
}
|
|
}
|
|
|
|
// 默认单窗口（全图）
|
|
if (!dynamic_roi_.enable && windows_.empty()) {
|
|
windows_.push_back({0, 0, 0, 0}); // 0表示全图
|
|
}
|
|
|
|
input_queue_ = ctx.input_queue;
|
|
output_queues_ = ctx.output_queues;
|
|
if (!input_queue_) {
|
|
LogError("[ai_shoe_det] no input queue");
|
|
return false;
|
|
}
|
|
|
|
infer_backend_ = ctx.infer_backend;
|
|
if (!infer_backend_) {
|
|
LogError("[ai_shoe_det] no infer backend");
|
|
return false;
|
|
}
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
std::string err;
|
|
model_handle_ = infer_backend_->LoadModel(model_path_, err);
|
|
if (model_handle_ == kInvalidModelHandle) {
|
|
LogError("[ai_shoe_det] failed to load model: " + err);
|
|
return false;
|
|
}
|
|
input_buf_.resize(model_w_ * model_h_ * 3);
|
|
LogInfo("[ai_shoe_det] model loaded: " + model_path_);
|
|
#else
|
|
LogWarn("[ai_shoe_det] RKNN disabled");
|
|
#endif
|
|
|
|
return true;
|
|
}
|
|
|
|
bool Start() override {
|
|
LogInfo("[ai_shoe_det] start mode=" + std::string(dynamic_roi_.enable ? "dynamic_roi" : "windows") +
|
|
" windows=" + std::to_string(windows_.size()) +
|
|
" append=" + std::string(append_detections_ ? "true" : "false") +
|
|
" infer_interval_ms=" + std::to_string(infer_interval_ms_));
|
|
return true;
|
|
}
|
|
|
|
void Stop() override {
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
if (model_handle_ != kInvalidModelHandle) {
|
|
infer_backend_->UnloadModel(model_handle_);
|
|
model_handle_ = kInvalidModelHandle;
|
|
}
|
|
#endif
|
|
LogInfo("[ai_shoe_det] stop");
|
|
}
|
|
|
|
NodeStatus Process(FramePtr frame) override {
|
|
if (!frame) return NodeStatus::DROP;
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
if (infer_interval_ms_ > 0 && frame->pts > 0) {
|
|
const int64_t pts_ms = static_cast<int64_t>(frame->pts / 1000ULL);
|
|
const int64_t delta_ms = pts_ms - last_infer_pts_ms_;
|
|
if (last_infer_pts_ms_ > 0 && delta_ms > 0 && delta_ms < infer_interval_ms_) {
|
|
Push(frame);
|
|
return NodeStatus::OK;
|
|
}
|
|
last_infer_pts_ms_ = pts_ms;
|
|
}
|
|
RunDetection(frame);
|
|
#endif
|
|
Push(frame);
|
|
return NodeStatus::OK;
|
|
}
|
|
|
|
private:
|
|
void Push(FramePtr frame) {
|
|
for (auto& q : output_queues_) q->Push(frame);
|
|
}
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
|
|
void RunDetection(FramePtr frame) {
|
|
if (!frame->data || frame->data_size == 0) return;
|
|
|
|
const int src_w = frame->width;
|
|
const int src_h = frame->height;
|
|
|
|
std::vector<DetBox> all_dets;
|
|
|
|
std::vector<DetWindow> active_windows = windows_;
|
|
if (dynamic_roi_.enable) {
|
|
active_windows = BuildDynamicWindows(frame, src_w, src_h);
|
|
if (active_windows.empty()) {
|
|
return;
|
|
}
|
|
}
|
|
|
|
// 对每个窗口进行检测
|
|
for (const auto& win : active_windows) {
|
|
auto dets = DetectWindow(frame, src_w, src_h, win);
|
|
all_dets.insert(all_dets.end(), dets.begin(), dets.end());
|
|
}
|
|
|
|
// NMS
|
|
all_dets = ApplyNMS(all_dets, nms_thresh_);
|
|
|
|
// 填充结果
|
|
if (!frame->det) {
|
|
frame->det = std::make_shared<DetectionResult>();
|
|
}
|
|
|
|
if (!append_detections_) {
|
|
frame->det->items.clear();
|
|
}
|
|
frame->det->img_w = src_w;
|
|
frame->det->img_h = src_h;
|
|
|
|
for (const auto& d : all_dets) {
|
|
Detection item;
|
|
item.bbox = {d.x, d.y, d.w, d.h};
|
|
item.score = d.conf;
|
|
item.cls_id = d.class_id;
|
|
frame->det->items.push_back(item);
|
|
}
|
|
}
|
|
|
|
std::vector<DetWindow> BuildDynamicWindows(FramePtr frame, int src_w, int src_h) const {
|
|
std::vector<DetWindow> rois;
|
|
if (!frame->det) return rois;
|
|
|
|
std::vector<Detection> persons;
|
|
persons.reserve(frame->det->items.size());
|
|
for (const auto& det : frame->det->items) {
|
|
if (det.cls_id != dynamic_roi_.person_class_id) continue;
|
|
if (det.bbox.h < dynamic_roi_.min_person_height) continue;
|
|
persons.push_back(det);
|
|
}
|
|
|
|
if (persons.empty()) return rois;
|
|
|
|
std::sort(persons.begin(), persons.end(),
|
|
[](const Detection& a, const Detection& b) { return a.score > b.score; });
|
|
|
|
const size_t max_rois = std::min(
|
|
persons.size(), static_cast<size_t>(std::max(1, dynamic_roi_.max_rois)));
|
|
rois.reserve(max_rois);
|
|
|
|
for (size_t i = 0; i < max_rois; ++i) {
|
|
const Rect& bbox = persons[i].bbox;
|
|
const float roi_x = bbox.x + bbox.w * dynamic_roi_.x_offset;
|
|
const float roi_y = bbox.y + bbox.h * dynamic_roi_.y_offset;
|
|
const float roi_w = bbox.w * dynamic_roi_.width_scale;
|
|
const float roi_h = bbox.h * dynamic_roi_.height_scale;
|
|
|
|
const int x0 = std::max(0, static_cast<int>(roi_x));
|
|
const int y0 = std::max(0, static_cast<int>(roi_y));
|
|
const int x1 = std::min(src_w, static_cast<int>(roi_x + roi_w));
|
|
const int y1 = std::min(src_h, static_cast<int>(roi_y + roi_h));
|
|
|
|
if (x1 <= x0 || y1 <= y0) continue;
|
|
rois.push_back({x0, y0, x1 - x0, y1 - y0});
|
|
}
|
|
|
|
return rois;
|
|
}
|
|
|
|
std::vector<DetBox> DetectWindow(FramePtr frame, int src_w, int src_h, const DetWindow& win) {
|
|
std::vector<DetBox> dets;
|
|
|
|
// 确定裁剪区域
|
|
int win_x, win_y, win_w, win_h;
|
|
if (win.w == 0 || win.h == 0) {
|
|
win_x = 0; win_y = 0; win_w = src_w; win_h = src_h;
|
|
} else {
|
|
win_x = std::max(0, std::min(win.x, src_w - 1));
|
|
win_y = std::max(0, std::min(win.y, src_h - 1));
|
|
win_w = std::min(win.w, src_w - win_x);
|
|
win_h = std::min(win.h, src_h - win_y);
|
|
}
|
|
|
|
if (win_w <= 0 || win_h <= 0) return dets;
|
|
|
|
// 获取源数据
|
|
const uint8_t* src = frame->planes[0].data ? frame->planes[0].data : frame->data;
|
|
int src_stride = frame->planes[0].stride > 0 ? frame->planes[0].stride
|
|
: (frame->stride > 0 ? frame->stride : src_w * 3);
|
|
|
|
if (!src || src_stride <= 0) return dets;
|
|
|
|
// 裁剪到临时缓冲区
|
|
std::vector<uint8_t> crop_buf(static_cast<size_t>(win_w) * win_h * 3);
|
|
for (int row = 0; row < win_h; ++row) {
|
|
const uint8_t* src_row = src + (win_y + row) * src_stride + win_x * 3;
|
|
uint8_t* dst_row = crop_buf.data() + row * win_w * 3;
|
|
memcpy(dst_row, src_row, static_cast<size_t>(win_w) * 3);
|
|
}
|
|
|
|
// Resize 到模型输入尺寸
|
|
ResizeRgbBilinear(crop_buf.data(), win_w, win_h, win_w * 3,
|
|
input_buf_.data(), model_w_, model_h_, model_w_ * 3);
|
|
|
|
// 推理
|
|
InferInput input;
|
|
input.width = model_w_;
|
|
input.height = model_h_;
|
|
input.is_nhwc = true;
|
|
input.data = input_buf_.data();
|
|
input.size = input_buf_.size();
|
|
input.type = RKNN_TENSOR_UINT8;
|
|
|
|
auto r = infer_backend_->InferBorrowed(model_handle_, input);
|
|
if (!r.success || r.outputs.empty() || !r.outputs[0].data) {
|
|
LogWarn("[ai_shoe_det] inference failed");
|
|
return dets;
|
|
}
|
|
|
|
// 解析输出
|
|
dets = ParseOutput(r.outputs, win_x, win_y, win_w, win_h);
|
|
return dets;
|
|
}
|
|
|
|
std::vector<DetBox> ParseOutput(const std::vector<AiScheduler::BorrowedOutput>& outputs,
|
|
int win_x, int win_y, int win_w, int win_h) {
|
|
std::vector<DetBox> dets;
|
|
if (outputs.empty() || !outputs[0].data) return dets;
|
|
|
|
if (outputs[0].type != RKNN_TENSOR_FLOAT32 || outputs[0].size < sizeof(float) * 5) {
|
|
LogWarn("[ai_shoe_det] unsupported output tensor type or size");
|
|
return dets;
|
|
}
|
|
|
|
// 假设 YOLOv8 输出格式: [num_boxes, 5] = x, y, w, h, conf
|
|
// 实际格式可能不同，需要根据模型调整
|
|
const float* data = reinterpret_cast<const float*>(outputs[0].data);
|
|
const size_t value_count = outputs[0].size / sizeof(float);
|
|
const int num_boxes = static_cast<int>(std::min<size_t>(8400, value_count / 5));
|
|
if (num_boxes <= 0) return dets;
|
|
|
|
float scale_x = static_cast<float>(win_w) / model_w_;
|
|
float scale_y = static_cast<float>(win_h) / model_h_;
|
|
|
|
for (int i = 0; i < num_boxes; ++i) {
|
|
float x = data[i * 5 + 0];
|
|
float y = data[i * 5 + 1];
|
|
float w = data[i * 5 + 2];
|
|
float h = data[i * 5 + 3];
|
|
float conf = data[i * 5 + 4];
|
|
|
|
if (conf < conf_thresh_) continue;
|
|
|
|
DetBox box;
|
|
box.x = win_x + x * scale_x;
|
|
box.y = win_y + y * scale_y;
|
|
box.w = w * scale_x;
|
|
box.h = h * scale_y;
|
|
box.conf = conf;
|
|
box.class_id = dynamic_roi_.shoe_class_id;
|
|
dets.push_back(box);
|
|
}
|
|
|
|
return dets;
|
|
}
|
|
|
|
std::vector<DetBox> ApplyNMS(std::vector<DetBox>& dets, float thresh) {
|
|
if (dets.empty()) return dets;
|
|
|
|
std::sort(dets.begin(), dets.end(),
|
|
[](const DetBox& a, const DetBox& b) { return a.conf > b.conf; });
|
|
|
|
std::vector<DetBox> keep;
|
|
std::vector<bool> suppressed(dets.size(), false);
|
|
|
|
for (size_t i = 0; i < dets.size(); ++i) {
|
|
if (suppressed[i]) continue;
|
|
keep.push_back(dets[i]);
|
|
|
|
for (size_t j = i + 1; j < dets.size(); ++j) {
|
|
if (suppressed[j]) continue;
|
|
float iou = ComputeIoU(dets[i], dets[j]);
|
|
if (iou > thresh) {
|
|
suppressed[j] = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
return keep;
|
|
}
|
|
|
|
float ComputeIoU(const DetBox& a, const DetBox& b) {
|
|
float x1 = std::max(a.x, b.x);
|
|
float y1 = std::max(a.y, b.y);
|
|
float x2 = std::min(a.x + a.w, b.x + b.w);
|
|
float y2 = std::min(a.y + a.h, b.y + b.h);
|
|
|
|
float inter = std::max(0.0f, x2 - x1) * std::max(0.0f, y2 - y1);
|
|
float area_a = a.w * a.h;
|
|
float area_b = b.w * b.h;
|
|
float uni = area_a + area_b - inter;
|
|
|
|
return uni > 0 ? inter / uni : 0;
|
|
}
|
|
|
|
void ResizeRgbBilinear(const uint8_t* src, int src_w, int src_h, int src_stride,
|
|
uint8_t* dst, int dst_w, int dst_h, int dst_stride) {
|
|
float scale_x = static_cast<float>(src_w) / dst_w;
|
|
float scale_y = static_cast<float>(src_h) / dst_h;
|
|
|
|
for (int y = 0; y < dst_h; ++y) {
|
|
float fy = y * scale_y;
|
|
int y0 = static_cast<int>(fy);
|
|
int y1 = std::min(y0 + 1, src_h - 1);
|
|
float dy = fy - y0;
|
|
|
|
for (int x = 0; x < dst_w; ++x) {
|
|
float fx = x * scale_x;
|
|
int x0 = static_cast<int>(fx);
|
|
int x1 = std::min(x0 + 1, src_w - 1);
|
|
float dx = fx - x0;
|
|
|
|
for (int c = 0; c < 3; ++c) {
|
|
float v00 = src[y0 * src_stride + x0 * 3 + c];
|
|
float v01 = src[y0 * src_stride + x1 * 3 + c];
|
|
float v10 = src[y1 * src_stride + x0 * 3 + c];
|
|
float v11 = src[y1 * src_stride + x1 * 3 + c];
|
|
|
|
float v = v00 * (1-dx) * (1-dy) +
|
|
v01 * dx * (1-dy) +
|
|
v10 * (1-dx) * dy +
|
|
v11 * dx * dy;
|
|
|
|
dst[y * dst_stride + x * 3 + c] = static_cast<uint8_t>(v);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#endif
|
|
|
|
std::string id_;
|
|
std::string model_path_;
|
|
int model_w_ = 640;
|
|
int model_h_ = 640;
|
|
float conf_thresh_ = 0.25f;
|
|
float nms_thresh_ = 0.45f;
|
|
bool append_detections_ = false;
|
|
int64_t infer_interval_ms_ = 0;
|
|
int64_t last_infer_pts_ms_ = 0;
|
|
DynamicRoiConfig dynamic_roi_;
|
|
std::vector<DetWindow> windows_;
|
|
std::vector<uint8_t> input_buf_;
|
|
|
|
std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
|
|
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
std::shared_ptr<IInferBackend> infer_backend_;
|
|
ModelHandle model_handle_ = kInvalidModelHandle;
|
|
#endif
|
|
};
|
|
|
|
REGISTER_NODE(AiShoeDetNode, "ai_shoe_det");
|
|
|
|
} // namespace rk3588
|