312 lines
10 KiB
C++
312 lines
10 KiB
C++
/**
|
|
* ai_scrfd_sliding - SCRFD with sliding window detection
|
|
*
|
|
* Features:
|
|
* 1. Resize input to target height (640) keeping approximate ratio
|
|
* 2. Split into multiple 640x640 windows
|
|
* 3. Detect on each window and merge results
|
|
*
|
|
* For 1080p: resize to 1280x640, 2 windows
|
|
* For 1440p: resize to 2560x640, 4 windows
|
|
*/
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "face/face_detection_utils.h"
|
|
#include "face/face_result.h"
|
|
#include "face/scrfd_detector.h"
|
|
#include "hw/i_infer_backend.h"
|
|
#include "node.h"
|
|
#include "utils/dma_alloc.h"
|
|
#include "utils/logger.h"
|
|
|
|
namespace rk3588 {
|
|
|
|
using namespace face_detection;
|
|
|
|
class AiScrfdSlidingNode : public INode {
|
|
public:
|
|
std::string Id() const override { return id_; }
|
|
std::string Type() const override { return "ai_scrfd_sliding"; }
|
|
|
|
bool Init(const SimpleJson& config, const NodeContext& ctx) override {
|
|
id_ = config.ValueOr<std::string>("id", "scrfd_sliding");
|
|
model_path_ = config.ValueOr<std::string>("model_path",
|
|
"./models/scrfd_500m_640.rknn");
|
|
|
|
// Detection parameters
|
|
det_cfg_.conf_thresh = config.ValueOr<float>("conf_thresh", 0.3f);
|
|
det_cfg_.nms_thresh = config.ValueOr<float>("nms_thresh", 0.4f);
|
|
det_cfg_.max_faces = config.ValueOr<int>("max_faces", 50);
|
|
det_cfg_.output_landmarks = config.ValueOr<bool>("output_landmarks", true);
|
|
|
|
model_w_ = 640;
|
|
model_h_ = 640;
|
|
|
|
// Initialize detector
|
|
detector_.Init(model_w_, model_h_);
|
|
|
|
// Parse sliding windows config
|
|
// If not configured, auto-calculate based on input resolution
|
|
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()) {
|
|
Window 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);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Target resize height (default 640)
|
|
target_height_ = config.ValueOr<int>("target_height", 640);
|
|
|
|
input_queue_ = ctx.input_queue;
|
|
output_queues_ = ctx.output_queues;
|
|
if (!input_queue_) {
|
|
LogError("[ai_scrfd_sliding] no input queue");
|
|
return false;
|
|
}
|
|
|
|
infer_backend_ = ctx.infer_backend;
|
|
if (!infer_backend_) {
|
|
LogError("[ai_scrfd_sliding] 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_scrfd_sliding] failed to load model: " + err);
|
|
return false;
|
|
}
|
|
|
|
input_buf_.resize(model_w_ * model_h_ * 3);
|
|
|
|
LogInfo("[ai_scrfd_sliding] model loaded: " + model_path_);
|
|
#else
|
|
LogWarn("[ai_scrfd_sliding] RKNN disabled");
|
|
#endif
|
|
|
|
return true;
|
|
}
|
|
|
|
bool Start() override {
|
|
LogInfo("[ai_scrfd_sliding] start, windows=" + std::to_string(windows_.size()));
|
|
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_scrfd_sliding] stop");
|
|
}
|
|
|
|
NodeStatus Process(FramePtr frame) override {
|
|
if (!frame) return NodeStatus::DROP;
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
RunDetection(frame);
|
|
#endif
|
|
|
|
Push(frame);
|
|
return NodeStatus::OK;
|
|
}
|
|
|
|
private:
|
|
struct Window {
|
|
int x, y, w, h;
|
|
};
|
|
|
|
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;
|
|
|
|
if (frame->DmaFd() >= 0) frame->SyncStart();
|
|
|
|
// Calculate windows if not pre-configured
|
|
std::vector<Window> windows = windows_;
|
|
if (windows.empty()) {
|
|
windows = CalculateWindows(src_w, src_h);
|
|
}
|
|
|
|
std::vector<FaceDetItem> all_detections;
|
|
|
|
const uint8_t* src = frame->planes[0].data ? frame->planes[0].data : frame->data;
|
|
const int src_stride = frame->planes[0].stride > 0 ? frame->planes[0].stride
|
|
: (frame->stride > 0 ? frame->stride : frame->width * 3);
|
|
|
|
// Process each window - crop from original, then resize to 640x640
|
|
for (size_t i = 0; i < windows.size(); ++i) {
|
|
const auto& win = windows[i];
|
|
auto dets = DetectWindowFromSource(src, src_w, src_h, src_stride, win);
|
|
|
|
// Detections are already in original coordinates
|
|
all_detections.insert(all_detections.end(), dets.begin(), dets.end());
|
|
}
|
|
|
|
// Apply NMS
|
|
all_detections = detector_.ApplyNMS(all_detections, det_cfg_.nms_thresh);
|
|
|
|
if (all_detections.size() > static_cast<size_t>(det_cfg_.max_faces)) {
|
|
all_detections.resize(det_cfg_.max_faces);
|
|
}
|
|
|
|
FaceDetResult result;
|
|
result.img_w = src_w;
|
|
result.img_h = src_h;
|
|
result.model_name = "scrfd_sliding";
|
|
result.faces = std::move(all_detections);
|
|
|
|
frame->face_det = std::make_shared<FaceDetResult>(std::move(result));
|
|
}
|
|
|
|
std::vector<Window> CalculateWindows(int src_w, int src_h) {
|
|
std::vector<Window> windows;
|
|
|
|
// Strategy: Split source image into overlapping 640x640 regions
|
|
// For 1080p: 1920x1080 -> 3x2 grid (6 windows)
|
|
// For 1440p: 2560x1440 -> 4x2 grid (8 windows)
|
|
|
|
// Calculate step size (with overlap)
|
|
int step_x = (src_w <= 640) ? src_w : (src_w - 640) / ((src_w + 639) / 640 - 1);
|
|
int step_y = (src_h <= 640) ? src_h : (src_h - 640) / ((src_h + 639) / 640 - 1);
|
|
|
|
if (step_x < 640) step_x = 640;
|
|
if (step_y < 640) step_y = 640;
|
|
|
|
for (int y = 0; y < src_h; y += step_y) {
|
|
for (int x = 0; x < src_w; x += step_x) {
|
|
Window win;
|
|
win.x = x;
|
|
win.y = y;
|
|
win.w = 640;
|
|
win.h = 640;
|
|
windows.push_back(win);
|
|
|
|
// Stop if we've covered the width
|
|
if (x + 640 >= src_w) break;
|
|
}
|
|
// Stop if we've covered the height
|
|
if (y + 640 >= src_h) break;
|
|
}
|
|
|
|
LogInfo("[ai_scrfd_sliding] Auto-calculated: " + std::to_string(windows.size()) + " windows for " + std::to_string(src_w) + "x" + std::to_string(src_h));
|
|
|
|
return windows;
|
|
}
|
|
|
|
std::vector<FaceDetItem> DetectWindowFromSource(const uint8_t* src, int src_w, int src_h, int src_stride, const Window& win) {
|
|
std::vector<FaceDetItem> dets;
|
|
|
|
// Clamp window to source bounds
|
|
int win_x = std::max(0, std::min(win.x, src_w - 1));
|
|
int win_y = std::max(0, std::min(win.y, src_h - 1));
|
|
int win_w = std::min(win.w, src_w - win_x);
|
|
int win_h = std::min(win.h, src_h - win_y);
|
|
|
|
if (win_w <= 0 || win_h <= 0) {
|
|
LogWarn("[ai_scrfd_sliding] Invalid window");
|
|
return dets;
|
|
}
|
|
|
|
// Crop from source
|
|
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 to 640x640
|
|
std::vector<uint8_t> model_input(640 * 640 * 3);
|
|
ResizeRgbBilinear(crop_buf.data(), win_w, win_h, win_w * 3,
|
|
model_input.data(), 640, 640, false);
|
|
|
|
// NPU inference
|
|
InferInput input;
|
|
input.width = 640;
|
|
input.height = 640;
|
|
input.is_nhwc = true;
|
|
input.data = model_input.data();
|
|
input.size = model_input.size();
|
|
input.type = RKNN_TENSOR_UINT8;
|
|
|
|
auto r = infer_backend_->InferBorrowed(model_handle_, input);
|
|
if (!r.success || r.outputs.empty()) {
|
|
LogWarn("[ai_scrfd_sliding] inference failed");
|
|
return dets;
|
|
}
|
|
|
|
// Decode (get detections in 640x640 coordinates)
|
|
dets = detector_.Decode(r.outputs, 640, 640, det_cfg_);
|
|
|
|
// Map back to original coordinates
|
|
float scale_x = static_cast<float>(win_w) / 640.0f;
|
|
float scale_y = static_cast<float>(win_h) / 640.0f;
|
|
|
|
for (auto& det : dets) {
|
|
det.bbox.x = win_x + det.bbox.x * scale_x;
|
|
det.bbox.y = win_y + det.bbox.y * scale_y;
|
|
det.bbox.w *= scale_x;
|
|
det.bbox.h *= scale_y;
|
|
if (det.has_landmarks) {
|
|
for (auto& lm : det.landmarks) {
|
|
lm.x = win_x + lm.x * scale_x;
|
|
lm.y = win_y + lm.y * scale_y;
|
|
}
|
|
}
|
|
}
|
|
|
|
return dets;
|
|
}
|
|
|
|
#endif
|
|
|
|
std::string id_;
|
|
std::string model_path_;
|
|
ScrfdConfig det_cfg_;
|
|
ScrfdDetector detector_;
|
|
int model_w_ = 640;
|
|
int model_h_ = 640;
|
|
int target_height_ = 640;
|
|
|
|
std::vector<Window> windows_;
|
|
|
|
std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
|
|
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;
|
|
std::shared_ptr<IInferBackend> infer_backend_;
|
|
|
|
#if defined(RK3588_ENABLE_RKNN)
|
|
ModelHandle model_handle_ = kInvalidModelHandle;
|
|
std::vector<uint8_t> input_buf_;
|
|
#endif
|
|
};
|
|
|
|
REGISTER_NODE(AiScrfdSlidingNode, "ai_scrfd_sliding");
|
|
|
|
} // namespace rk3588
|