优化性能,对齐prd
This commit is contained in:
parent
bb22c34696
commit
f8aa377a79
39
Readme.md
39
Readme.md
@ -609,6 +609,7 @@ alarm 插件采用**模块化 Actions 架构**,支持多种报警动作的灵
|
||||
│ └──────────────┘ │ AlarmActions (可配置) │ │
|
||||
│ │ - LogAction → 日志输出 │ │
|
||||
│ │ - HttpAction → HTTP 回调 │ │
|
||||
│ │ - ExternalApiAction → 外部API报警 │ │
|
||||
│ │ - SnapshotAction → 截图+上传 │ │
|
||||
│ │ - ClipAction → 视频片段+上传 │ │
|
||||
│ └──────────────────────────────────┘ │
|
||||
@ -641,6 +642,18 @@ alarm 插件采用**模块化 Actions 架构**,支持多种报警动作的灵
|
||||
"timeout_ms": 3000,
|
||||
"include_media_url": true
|
||||
},
|
||||
"external_api": {
|
||||
"enable": false,
|
||||
"getTokenUrl": "http://业务系统/api/getToken",
|
||||
"putMessageUrl": "http://业务系统/api/putMessage",
|
||||
"tenantCode": "32",
|
||||
"channelNo": "${vod_channelNo}",
|
||||
"timeout_ms": 3000,
|
||||
"include_media_url": true,
|
||||
"token_header": "X-Access-Token",
|
||||
"token_json_path": "responseBody.token",
|
||||
"token_cache_sec": 600
|
||||
},
|
||||
"snapshot": {
|
||||
"enable": true,
|
||||
"format": "jpg",
|
||||
@ -664,6 +677,32 @@ alarm 插件采用**模块化 Actions 架构**,支持多种报警动作的灵
|
||||
}
|
||||
```
|
||||
|
||||
**external_api 动作说明(对接外部系统报警:先取 token,再上报告警):**
|
||||
|
||||
- 取 token:对 `getTokenUrl` 发起 HTTP POST,默认从响应 JSON 的 `responseBody.token` 读取(可用 `token_json_path` 修改点号路径)。
|
||||
- 上报告警:对 `putMessageUrl` 发起 HTTP POST,并携带 Header:`Content-Type: application/json`、`X-Access-Token: <token>`(可用 `token_header` 修改)。
|
||||
- 上报请求体(固定结构):
|
||||
|
||||
```json
|
||||
{
|
||||
"tenantCode": "32",
|
||||
"channelNo": "<vod_channelNo>",
|
||||
"alarmContent": "<msg>",
|
||||
"alarmTime": "YYYY-MM-DD HH:MM:SS",
|
||||
"picInfo": [{"url": "<picUrl>"}],
|
||||
"videoInfo": [{"url": "<videoUrl>"}]
|
||||
}
|
||||
```
|
||||
|
||||
- `alarmTime` 使用本地时间格式化。
|
||||
- `picUrl/videoUrl` 来自 `snapshot/clip` 动作上传后的 URL(通常为 MinIO 返回链接);若为空则对应数组为空。
|
||||
- `alarmContent`:若配置中未提供 `alarmContent`,默认使用 `rule_name`。
|
||||
|
||||
**external_api 参数说明:**
|
||||
- 必填:`getTokenUrl`, `putMessageUrl`, `tenantCode`, `channelNo`
|
||||
- 常用可选:`timeout_ms`(默认 3000)、`include_media_url`(默认 true)、`token_cache_sec`(默认 600;为 0 时仅在 401/403 时刷新 token)
|
||||
- 其它可选:`token_header`、`token_json_path`、`alarmContent`、`max_queue_size` / `queue_policy`、`max_retries` / `retry_backoff_ms`
|
||||
|
||||
**规则参数说明:**
|
||||
- `class_ids`: 要监控的类别 ID 数组(与模型输出对应)
|
||||
- `objects`: 要监控的类别名称数组(需配合 `labels` 使用)
|
||||
|
||||
@ -3,7 +3,7 @@
|
||||
"metrics_port": 9000,
|
||||
"web_root": "web"
|
||||
},
|
||||
"queue": { "size": 32, "strategy": "drop_oldest" },
|
||||
"queue": { "size": 4, "strategy": "drop_oldest" },
|
||||
|
||||
"templates": {
|
||||
"strict_minio_alarm_pipeline": {
|
||||
@ -63,6 +63,7 @@
|
||||
"type": "preprocess",
|
||||
"role": "filter",
|
||||
"enable": true,
|
||||
"queue": { "size": 2, "policy": "drop_oldest" },
|
||||
"dst_w": "${src_w}",
|
||||
"dst_h": "${src_h}",
|
||||
"dst_format": "nv12",
|
||||
@ -81,14 +82,15 @@
|
||||
"name": "person_in_view",
|
||||
"class_ids": [0],
|
||||
"roi": { "x": 0.0, "y": 0.0, "w": 1.0, "h": 1.0 },
|
||||
"min_duration_ms": 0,
|
||||
"cooldown_ms": 3000
|
||||
"min_duration_ms": 500,
|
||||
"cooldown_ms": 15000
|
||||
}
|
||||
],
|
||||
"actions": {
|
||||
"log": { "enable": true, "level": "info" },
|
||||
"snapshot": {
|
||||
"enable": true,
|
||||
"min_interval_ms": 15000,
|
||||
"format": "jpg",
|
||||
"quality": 85,
|
||||
"upload": {
|
||||
@ -102,6 +104,7 @@
|
||||
},
|
||||
"clip": {
|
||||
"enable": true,
|
||||
"min_interval_ms": 15000,
|
||||
"pre_sec": 5,
|
||||
"post_sec": 10,
|
||||
"format": "mp4",
|
||||
@ -129,6 +132,7 @@
|
||||
"type": "publish",
|
||||
"role": "filter",
|
||||
"enable": true,
|
||||
"queue": { "size": 2, "policy": "drop_oldest" },
|
||||
"codec": "h264",
|
||||
"fps": "${fps}",
|
||||
"gop": "${gop}",
|
||||
|
||||
@ -76,6 +76,40 @@ public:
|
||||
// Synchronous inference
|
||||
InferResult Infer(ModelHandle handle, const InferInput& input);
|
||||
|
||||
struct BorrowedOutput {
|
||||
const uint8_t* data = nullptr;
|
||||
size_t size = 0;
|
||||
int index = 0;
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
rknn_tensor_type type = RKNN_TENSOR_UINT8;
|
||||
int32_t zp = 0;
|
||||
float scale = 1.0f;
|
||||
std::vector<uint32_t> dims;
|
||||
#endif
|
||||
};
|
||||
|
||||
// Borrowed inference avoids per-call output allocations/copies by using per-model preallocated buffers.
|
||||
// The returned result holds the per-model inference lock; it must be destroyed before the next InferBorrowed()
|
||||
// on the same model handle.
|
||||
struct BorrowedInferResult {
|
||||
bool success = false;
|
||||
std::string error;
|
||||
std::vector<BorrowedOutput> outputs;
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
std::shared_ptr<void> keepalive; // keeps ModelContext alive
|
||||
std::unique_lock<std::mutex> infer_lock; // holds ModelContext::infer_mutex
|
||||
#endif
|
||||
|
||||
BorrowedInferResult() = default;
|
||||
BorrowedInferResult(BorrowedInferResult&&) noexcept = default;
|
||||
BorrowedInferResult& operator=(BorrowedInferResult&&) noexcept = default;
|
||||
BorrowedInferResult(const BorrowedInferResult&) = delete;
|
||||
BorrowedInferResult& operator=(const BorrowedInferResult&) = delete;
|
||||
};
|
||||
|
||||
BorrowedInferResult InferBorrowed(ModelHandle handle, const InferInput& input);
|
||||
|
||||
// Async inference (submits to internal queue, calls callback when done)
|
||||
// For now, this is a simple wrapper around sync Infer
|
||||
void InferAsync(ModelHandle handle, const InferInput& input, InferCallback callback);
|
||||
@ -97,6 +131,7 @@ private:
|
||||
std::vector<uint8_t> model_data;
|
||||
std::vector<rknn_tensor_attr> input_attrs;
|
||||
std::vector<rknn_tensor_attr> output_attrs;
|
||||
std::vector<std::vector<uint8_t>> output_buffers; // preallocated output buffers
|
||||
uint32_t n_input = 0;
|
||||
uint32_t n_output = 0;
|
||||
int input_w = 0;
|
||||
@ -113,7 +148,8 @@ private:
|
||||
}
|
||||
};
|
||||
|
||||
std::unordered_map<ModelHandle, std::shared_ptr<ModelContext>> models_;
|
||||
std::unordered_map<ModelHandle, std::shared_ptr<ModelContext>> models_by_handle_;
|
||||
std::unordered_map<std::string, std::weak_ptr<ModelContext>> models_by_path_;
|
||||
#endif
|
||||
|
||||
mutable std::mutex models_mutex_; // Protects models_ map
|
||||
|
||||
@ -73,6 +73,19 @@ public:
|
||||
return true;
|
||||
}
|
||||
|
||||
// Blocks until an item is available or the queue is stopped.
|
||||
bool Pop(T& out) {
|
||||
std::unique_lock<std::mutex> lock(mu_);
|
||||
data_cv_.wait(lock, [&] { return size_ > 0 || stop_; });
|
||||
if (size_ == 0) return false;
|
||||
out = std::move(buf_[head_]);
|
||||
head_ = (head_ + 1) % capacity_;
|
||||
--size_;
|
||||
++popped_;
|
||||
space_cv_.notify_one();
|
||||
return true;
|
||||
}
|
||||
|
||||
void Stop() {
|
||||
std::lock_guard<std::mutex> lock(mu_);
|
||||
stop_ = true;
|
||||
|
||||
@ -437,13 +437,131 @@ private:
|
||||
input.height = h;
|
||||
input.is_nhwc = true;
|
||||
|
||||
InferResult result = AiScheduler::Instance().Infer(model_handle_, input);
|
||||
auto result = AiScheduler::Instance().InferBorrowed(model_handle_, input);
|
||||
if (!result.success) {
|
||||
std::cerr << "[ai_yolo] inference failed: " << result.error << "\n";
|
||||
return;
|
||||
}
|
||||
|
||||
PostProcess(result.outputs, frame);
|
||||
PostProcessBorrowed(result.outputs, frame);
|
||||
}
|
||||
|
||||
void PostProcessBorrowed(const std::vector<AiScheduler::BorrowedOutput>& outputs, FramePtr frame) {
|
||||
std::vector<float> boxes;
|
||||
std::vector<float> obj_probs;
|
||||
std::vector<int> class_ids;
|
||||
int valid_count = 0;
|
||||
|
||||
if (yolo_version_ == YoloVersion::V5) {
|
||||
if (outputs.size() < 3) return;
|
||||
if (!outputs[0].data || !outputs[1].data || !outputs[2].data) return;
|
||||
|
||||
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 = ProcessFeatureMapV5(reinterpret_cast<int8_t*>(const_cast<uint8_t*>(outputs[0].data)), kAnchor0,
|
||||
grid_h0, grid_w0, model_input_h_, model_input_w_, stride0,
|
||||
boxes, obj_probs, class_ids, conf_thresh_,
|
||||
outputs[0].zp, outputs[0].scale);
|
||||
int cnt1 = ProcessFeatureMapV5(reinterpret_cast<int8_t*>(const_cast<uint8_t*>(outputs[1].data)), kAnchor1,
|
||||
grid_h1, grid_w1, model_input_h_, model_input_w_, stride1,
|
||||
boxes, obj_probs, class_ids, conf_thresh_,
|
||||
outputs[1].zp, outputs[1].scale);
|
||||
int cnt2 = ProcessFeatureMapV5(reinterpret_cast<int8_t*>(const_cast<uint8_t*>(outputs[2].data)), kAnchor2,
|
||||
grid_h2, grid_w2, model_input_h_, model_input_w_, stride2,
|
||||
boxes, obj_probs, class_ids, conf_thresh_,
|
||||
outputs[2].zp, outputs[2].scale);
|
||||
valid_count = cnt0 + cnt1 + cnt2;
|
||||
} else {
|
||||
if (outputs.empty()) return;
|
||||
if (!outputs[0].data || outputs[0].size == 0) return;
|
||||
|
||||
int num_boxes = 0;
|
||||
int num_channels = 4 + num_classes_;
|
||||
|
||||
if (outputs[0].dims.size() >= 3) {
|
||||
if (outputs[0].dims[1] == static_cast<uint32_t>(num_channels)) {
|
||||
num_boxes = static_cast<int>(outputs[0].dims[2]);
|
||||
} else if (outputs[0].dims[2] == static_cast<uint32_t>(num_channels)) {
|
||||
num_boxes = static_cast<int>(outputs[0].dims[1]);
|
||||
} else {
|
||||
num_boxes = 8400;
|
||||
}
|
||||
} else {
|
||||
num_boxes = static_cast<int>(outputs[0].size) / num_channels;
|
||||
}
|
||||
|
||||
if (outputs[0].type == RKNN_TENSOR_FLOAT32 ||
|
||||
outputs[0].type == RKNN_TENSOR_FLOAT16) {
|
||||
valid_count = ProcessOutputV8(reinterpret_cast<float*>(const_cast<uint8_t*>(outputs[0].data)),
|
||||
num_boxes, num_classes_,
|
||||
model_input_h_, model_input_w_,
|
||||
boxes, obj_probs, class_ids, conf_thresh_);
|
||||
} else {
|
||||
valid_count = ProcessOutputV8Int8(reinterpret_cast<int8_t*>(const_cast<uint8_t*>(outputs[0].data)),
|
||||
num_boxes, num_classes_,
|
||||
model_input_h_, model_input_w_,
|
||||
boxes, obj_probs, class_ids, conf_thresh_,
|
||||
outputs[0].zp, outputs[0].scale);
|
||||
}
|
||||
}
|
||||
|
||||
if (valid_count <= 0) return;
|
||||
|
||||
std::vector<int> indices(valid_count);
|
||||
for (int i = 0; i < valid_count; ++i) indices[i] = i;
|
||||
|
||||
QuickSortDescending(obj_probs, 0, valid_count - 1, indices);
|
||||
|
||||
std::set<int> 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<float>(model_input_w_) / frame->width;
|
||||
float scale_h = static_cast<float>(model_input_h_) / frame->height;
|
||||
|
||||
auto det_result = std::make_shared<DetectionResult>();
|
||||
det_result->img_w = frame->width;
|
||||
det_result->img_h = frame->height;
|
||||
det_result->model_name = (yolo_version_ == YoloVersion::V5) ? "yolov5" : "yolov8";
|
||||
|
||||
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 = static_cast<float>(Clamp(static_cast<int>(x1 / scale_w), 0, frame->width));
|
||||
det.bbox.y = static_cast<float>(Clamp(static_cast<int>(y1 / scale_h), 0, frame->height));
|
||||
det.bbox.w = static_cast<float>(Clamp(static_cast<int>(w / scale_w), 0, frame->width - static_cast<int>(det.bbox.x)));
|
||||
det.bbox.h = static_cast<float>(Clamp(static_cast<int>(h / scale_h), 0, frame->height - static_cast<int>(det.bbox.y)));
|
||||
det.track_id = -1;
|
||||
|
||||
if (det_result->items.size() < 3 && processed_ < 10) {
|
||||
std::cout << "[ai_yolo] det: raw(" << x1 << "," << y1 << "," << w << "," << h
|
||||
<< ") -> bbox(" << det.bbox.x << "," << det.bbox.y << ","
|
||||
<< det.bbox.w << "," << det.bbox.h << ") cls=" << cls_id
|
||||
<< " score=" << det.score << "\n";
|
||||
}
|
||||
|
||||
det_result->items.push_back(det);
|
||||
}
|
||||
|
||||
frame->det = det_result;
|
||||
}
|
||||
|
||||
void PostProcess(std::vector<InferOutput>& outputs, FramePtr frame) {
|
||||
|
||||
@ -21,6 +21,68 @@
|
||||
|
||||
namespace rk3588 {
|
||||
|
||||
namespace {
|
||||
|
||||
uint64_t NowEpochMs() {
|
||||
using namespace std::chrono;
|
||||
return static_cast<uint64_t>(duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count());
|
||||
}
|
||||
|
||||
class RateLimitedAction final : public IAlarmAction {
|
||||
public:
|
||||
RateLimitedAction(std::unique_ptr<IAlarmAction> inner, uint64_t min_interval_ms)
|
||||
: inner_(std::move(inner)), min_interval_ms_(min_interval_ms) {}
|
||||
|
||||
bool Init(const SimpleJson& config) override {
|
||||
// Inner is expected to be initialized already; keep compatibility if called.
|
||||
return inner_ ? inner_->Init(config) : false;
|
||||
}
|
||||
|
||||
void Execute(AlarmEvent& event, std::shared_ptr<Frame> frame) override {
|
||||
if (!inner_) return;
|
||||
if (min_interval_ms_ == 0) {
|
||||
inner_->Execute(event, std::move(frame));
|
||||
return;
|
||||
}
|
||||
|
||||
const uint64_t now = NowEpochMs();
|
||||
uint64_t last = last_exec_ms_.load(std::memory_order_relaxed);
|
||||
if (last != 0 && now > last && (now - last) < min_interval_ms_) {
|
||||
return;
|
||||
}
|
||||
// Best-effort: only one thread should update; compare-exchange avoids lost updates.
|
||||
last_exec_ms_.store(now, std::memory_order_relaxed);
|
||||
inner_->Execute(event, std::move(frame));
|
||||
}
|
||||
|
||||
void Drain() override {
|
||||
if (inner_) inner_->Drain();
|
||||
}
|
||||
|
||||
void Stop() override {
|
||||
if (inner_) inner_->Stop();
|
||||
}
|
||||
|
||||
std::string Name() const override {
|
||||
return inner_ ? inner_->Name() : std::string("unknown");
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<IAlarmAction> inner_;
|
||||
uint64_t min_interval_ms_ = 0;
|
||||
std::atomic<uint64_t> last_exec_ms_{0};
|
||||
};
|
||||
|
||||
std::unique_ptr<IAlarmAction> WrapIfRateLimited(std::unique_ptr<IAlarmAction> action, const SimpleJson& cfg) {
|
||||
if (!action) return nullptr;
|
||||
const uint64_t min_interval_ms = static_cast<uint64_t>(std::max<int64_t>(
|
||||
0, static_cast<int64_t>(cfg.ValueOr<int>("min_interval_ms", 0))));
|
||||
if (min_interval_ms == 0) return action;
|
||||
return std::make_unique<RateLimitedAction>(std::move(action), min_interval_ms);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
class AlarmNode : public INode {
|
||||
private:
|
||||
struct AlarmJob {
|
||||
@ -105,7 +167,7 @@ public:
|
||||
if (log_cfg->ValueOr<bool>("enable", true)) {
|
||||
auto action = std::make_unique<LogAction>();
|
||||
if (action->Init(*log_cfg)) {
|
||||
actions_.push_back(std::move(action));
|
||||
actions_.push_back(WrapIfRateLimited(std::move(action), *log_cfg));
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -115,7 +177,7 @@ public:
|
||||
if (snap_cfg->ValueOr<bool>("enable", false)) {
|
||||
auto action = std::make_unique<SnapshotAction>();
|
||||
if (action->Init(*snap_cfg)) {
|
||||
actions_.push_back(std::move(action));
|
||||
actions_.push_back(WrapIfRateLimited(std::move(action), *snap_cfg));
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -127,7 +189,7 @@ public:
|
||||
if (action->Init(*clip_cfg)) {
|
||||
auto* clip_ptr = static_cast<ClipAction*>(action.get());
|
||||
clip_ptr->SetPacketBuffer(packet_buffer_);
|
||||
actions_.push_back(std::move(action));
|
||||
actions_.push_back(WrapIfRateLimited(std::move(action), *clip_cfg));
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -137,7 +199,7 @@ public:
|
||||
if (http_cfg->ValueOr<bool>("enable", false)) {
|
||||
auto action = std::make_unique<HttpAction>();
|
||||
if (action->Init(*http_cfg)) {
|
||||
actions_.push_back(std::move(action));
|
||||
actions_.push_back(WrapIfRateLimited(std::move(action), *http_cfg));
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -147,7 +209,7 @@ public:
|
||||
if (ext_cfg->ValueOr<bool>("enable", false)) {
|
||||
auto action = std::make_unique<ExternalApiAction>();
|
||||
if (action->Init(*ext_cfg)) {
|
||||
actions_.push_back(std::move(action));
|
||||
actions_.push_back(WrapIfRateLimited(std::move(action), *ext_cfg));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -7,6 +7,32 @@
|
||||
|
||||
namespace rk3588 {
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
namespace {
|
||||
|
||||
uint32_t TensorTypeSizeBytes(rknn_tensor_type t) {
|
||||
switch (t) {
|
||||
case RKNN_TENSOR_INT8:
|
||||
case RKNN_TENSOR_UINT8:
|
||||
return 1;
|
||||
case RKNN_TENSOR_INT16:
|
||||
case RKNN_TENSOR_UINT16:
|
||||
case RKNN_TENSOR_FLOAT16:
|
||||
return 2;
|
||||
case RKNN_TENSOR_INT32:
|
||||
case RKNN_TENSOR_UINT32:
|
||||
case RKNN_TENSOR_FLOAT32:
|
||||
return 4;
|
||||
case RKNN_TENSOR_INT64:
|
||||
return 8;
|
||||
default:
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
#endif
|
||||
|
||||
AiScheduler& AiScheduler::Instance() {
|
||||
static AiScheduler instance;
|
||||
return instance;
|
||||
@ -24,7 +50,8 @@ void AiScheduler::Shutdown() {
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(models_mutex_);
|
||||
models_.clear();
|
||||
models_by_handle_.clear();
|
||||
models_by_path_.clear();
|
||||
}
|
||||
LogInfo("[AiScheduler] shutdown, total inferences: " + std::to_string(total_inferences_.load()) +
|
||||
", errors: " + std::to_string(total_errors_.load()));
|
||||
@ -33,6 +60,19 @@ void AiScheduler::Shutdown() {
|
||||
|
||||
ModelHandle AiScheduler::LoadModel(const std::string& model_path, std::string& err) {
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(models_mutex_);
|
||||
auto it = models_by_path_.find(model_path);
|
||||
if (it != models_by_path_.end()) {
|
||||
if (auto existing = it->second.lock()) {
|
||||
ModelHandle handle = next_handle_.fetch_add(1);
|
||||
models_by_handle_[handle] = existing;
|
||||
LogInfo("[AiScheduler] reused model: " + model_path + " (handle=" + std::to_string(handle) + ")");
|
||||
return handle;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Read model file
|
||||
std::ifstream file(model_path, std::ios::binary | std::ios::ate);
|
||||
if (!file.is_open()) {
|
||||
@ -86,6 +126,23 @@ ModelHandle AiScheduler::LoadModel(const std::string& model_path, std::string& e
|
||||
rknn_query(ctx->ctx, RKNN_QUERY_OUTPUT_ATTR, &ctx->output_attrs[i], sizeof(rknn_tensor_attr));
|
||||
}
|
||||
|
||||
// Preallocate output buffers for borrowed inference.
|
||||
ctx->output_buffers.resize(ctx->n_output);
|
||||
for (uint32_t i = 0; i < ctx->n_output; ++i) {
|
||||
uint32_t out_sz = ctx->output_attrs[i].size;
|
||||
if (out_sz == 0 && ctx->output_attrs[i].size_with_stride > 0) {
|
||||
out_sz = ctx->output_attrs[i].size_with_stride;
|
||||
}
|
||||
if (out_sz == 0 && ctx->output_attrs[i].n_elems > 0) {
|
||||
out_sz = ctx->output_attrs[i].n_elems * TensorTypeSizeBytes(ctx->output_attrs[i].type);
|
||||
}
|
||||
if (out_sz > 0) {
|
||||
ctx->output_buffers[i].resize(out_sz);
|
||||
} else {
|
||||
ctx->output_buffers[i].clear();
|
||||
}
|
||||
}
|
||||
|
||||
// Extract input dimensions
|
||||
if (ctx->input_attrs[0].fmt == RKNN_TENSOR_NCHW) {
|
||||
ctx->input_c = ctx->input_attrs[0].dims[1];
|
||||
@ -102,7 +159,8 @@ ModelHandle AiScheduler::LoadModel(const std::string& model_path, std::string& e
|
||||
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(models_mutex_);
|
||||
models_[handle] = ctx;
|
||||
models_by_handle_[handle] = ctx;
|
||||
models_by_path_[model_path] = ctx;
|
||||
}
|
||||
|
||||
LogInfo("[AiScheduler] loaded model: " + model_path +
|
||||
@ -124,9 +182,9 @@ void AiScheduler::UnloadModel(ModelHandle handle) {
|
||||
bool erased = false;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(models_mutex_);
|
||||
auto it = models_.find(handle);
|
||||
if (it != models_.end()) {
|
||||
models_.erase(it);
|
||||
auto it = models_by_handle_.find(handle);
|
||||
if (it != models_by_handle_.end()) {
|
||||
models_by_handle_.erase(it);
|
||||
erased = true;
|
||||
}
|
||||
}
|
||||
@ -143,8 +201,8 @@ bool AiScheduler::GetModelInfo(ModelHandle handle, ModelInfo& info) const {
|
||||
std::shared_ptr<ModelContext> ctx;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(models_mutex_);
|
||||
auto it = models_.find(handle);
|
||||
if (it == models_.end() || !it->second) {
|
||||
auto it = models_by_handle_.find(handle);
|
||||
if (it == models_by_handle_.end() || !it->second) {
|
||||
return false;
|
||||
}
|
||||
ctx = it->second;
|
||||
@ -171,8 +229,8 @@ InferResult AiScheduler::Infer(ModelHandle handle, const InferInput& input) {
|
||||
std::shared_ptr<ModelContext> ctx;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(models_mutex_);
|
||||
auto it = models_.find(handle);
|
||||
if (it == models_.end() || !it->second) {
|
||||
auto it = models_by_handle_.find(handle);
|
||||
if (it == models_by_handle_.end() || !it->second) {
|
||||
result.error = "Invalid model handle";
|
||||
total_errors_.fetch_add(1);
|
||||
return result;
|
||||
@ -263,6 +321,108 @@ InferResult AiScheduler::Infer(ModelHandle handle, const InferInput& input) {
|
||||
return result;
|
||||
}
|
||||
|
||||
AiScheduler::BorrowedInferResult AiScheduler::InferBorrowed(ModelHandle handle, const InferInput& input) {
|
||||
BorrowedInferResult result;
|
||||
|
||||
#if defined(RK3588_ENABLE_RKNN)
|
||||
std::shared_ptr<ModelContext> ctx;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(models_mutex_);
|
||||
auto it = models_by_handle_.find(handle);
|
||||
if (it == models_by_handle_.end() || !it->second) {
|
||||
result.error = "Invalid model handle";
|
||||
total_errors_.fetch_add(1);
|
||||
return result;
|
||||
}
|
||||
ctx = it->second;
|
||||
}
|
||||
|
||||
if (!input.data || input.size == 0) {
|
||||
result.error = "Invalid input data";
|
||||
total_errors_.fetch_add(1);
|
||||
return result;
|
||||
}
|
||||
|
||||
// Hold per-model inference lock for the lifetime of this result.
|
||||
result.infer_lock = std::unique_lock<std::mutex>(ctx->infer_mutex);
|
||||
result.keepalive = ctx;
|
||||
|
||||
// Setup input
|
||||
rknn_input inputs[1];
|
||||
memset(inputs, 0, sizeof(inputs));
|
||||
inputs[0].index = 0;
|
||||
inputs[0].type = RKNN_TENSOR_UINT8;
|
||||
inputs[0].size = input.size;
|
||||
inputs[0].fmt = input.is_nhwc ? RKNN_TENSOR_NHWC : RKNN_TENSOR_NCHW;
|
||||
inputs[0].buf = const_cast<void*>(input.data);
|
||||
inputs[0].pass_through = 0;
|
||||
|
||||
int ret = rknn_inputs_set(ctx->ctx, ctx->n_input, inputs);
|
||||
if (ret < 0) {
|
||||
result.error = "rknn_inputs_set failed: " + std::to_string(ret);
|
||||
total_errors_.fetch_add(1);
|
||||
return result;
|
||||
}
|
||||
|
||||
ret = rknn_run(ctx->ctx, nullptr);
|
||||
if (ret < 0) {
|
||||
result.error = "rknn_run failed: " + std::to_string(ret);
|
||||
total_errors_.fetch_add(1);
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<rknn_output> outputs(ctx->n_output);
|
||||
memset(outputs.data(), 0, sizeof(rknn_output) * ctx->n_output);
|
||||
for (uint32_t i = 0; i < ctx->n_output; ++i) {
|
||||
outputs[i].want_float = 0;
|
||||
outputs[i].index = i;
|
||||
if (i < ctx->output_buffers.size() && !ctx->output_buffers[i].empty()) {
|
||||
outputs[i].is_prealloc = 1;
|
||||
outputs[i].buf = ctx->output_buffers[i].data();
|
||||
outputs[i].size = static_cast<uint32_t>(ctx->output_buffers[i].size());
|
||||
} else {
|
||||
outputs[i].is_prealloc = 0;
|
||||
outputs[i].buf = nullptr;
|
||||
outputs[i].size = 0;
|
||||
}
|
||||
}
|
||||
|
||||
ret = rknn_outputs_get(ctx->ctx, ctx->n_output, outputs.data(), nullptr);
|
||||
if (ret < 0) {
|
||||
result.error = "rknn_outputs_get failed: " + std::to_string(ret);
|
||||
total_errors_.fetch_add(1);
|
||||
return result;
|
||||
}
|
||||
|
||||
result.outputs.resize(ctx->n_output);
|
||||
for (uint32_t i = 0; i < ctx->n_output; ++i) {
|
||||
auto& out = result.outputs[i];
|
||||
out.index = static_cast<int>(i);
|
||||
out.size = outputs[i].size;
|
||||
out.data = reinterpret_cast<const uint8_t*>(outputs[i].buf);
|
||||
out.type = ctx->output_attrs[i].type;
|
||||
out.zp = ctx->output_attrs[i].zp;
|
||||
out.scale = ctx->output_attrs[i].scale;
|
||||
out.dims.resize(ctx->output_attrs[i].n_dims);
|
||||
for (uint32_t d = 0; d < ctx->output_attrs[i].n_dims; ++d) {
|
||||
out.dims[d] = ctx->output_attrs[i].dims[d];
|
||||
}
|
||||
}
|
||||
|
||||
rknn_outputs_release(ctx->ctx, ctx->n_output, outputs.data());
|
||||
|
||||
result.success = true;
|
||||
total_inferences_.fetch_add(1);
|
||||
return result;
|
||||
|
||||
#else
|
||||
(void)handle;
|
||||
(void)input;
|
||||
result.error = "RKNN not enabled";
|
||||
return result;
|
||||
#endif
|
||||
}
|
||||
|
||||
void AiScheduler::InferAsync(ModelHandle handle, const InferInput& input, InferCallback callback) {
|
||||
// Simple implementation: just call sync Infer
|
||||
// Future: use a thread pool for true async
|
||||
|
||||
@ -526,39 +526,28 @@ bool Graph::Start() {
|
||||
}
|
||||
}
|
||||
FramePtr frame;
|
||||
while (true) {
|
||||
if (entry.context.input_queue->Pop(frame, std::chrono::milliseconds(100))) {
|
||||
if (!frame) {
|
||||
continue;
|
||||
}
|
||||
while (entry.context.input_queue->Pop(frame)) {
|
||||
if (!frame) continue;
|
||||
|
||||
const auto t0 = std::chrono::steady_clock::now();
|
||||
NodeStatus st = entry.node->Process(frame);
|
||||
const auto t1 = std::chrono::steady_clock::now();
|
||||
const auto ns = std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0).count();
|
||||
if (entry.metrics) {
|
||||
if (ns > 0) {
|
||||
entry.metrics->process_time_ns_total.fetch_add(static_cast<uint64_t>(ns),
|
||||
std::memory_order_relaxed);
|
||||
}
|
||||
if (st == NodeStatus::OK) {
|
||||
entry.metrics->ok_total.fetch_add(1, std::memory_order_relaxed);
|
||||
} else if (st == NodeStatus::DROP) {
|
||||
entry.metrics->drop_total.fetch_add(1, std::memory_order_relaxed);
|
||||
} else {
|
||||
entry.metrics->error_total.fetch_add(1, std::memory_order_relaxed);
|
||||
}
|
||||
const auto t0 = std::chrono::steady_clock::now();
|
||||
NodeStatus st = entry.node->Process(frame);
|
||||
const auto t1 = std::chrono::steady_clock::now();
|
||||
const auto ns = std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0).count();
|
||||
if (entry.metrics) {
|
||||
if (ns > 0) {
|
||||
entry.metrics->process_time_ns_total.fetch_add(static_cast<uint64_t>(ns),
|
||||
std::memory_order_relaxed);
|
||||
}
|
||||
if (st == NodeStatus::OK) {
|
||||
entry.metrics->ok_total.fetch_add(1, std::memory_order_relaxed);
|
||||
} else if (st == NodeStatus::DROP) {
|
||||
entry.metrics->drop_total.fetch_add(1, std::memory_order_relaxed);
|
||||
} else {
|
||||
entry.metrics->error_total.fetch_add(1, std::memory_order_relaxed);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Avoid busy-spin when upstream queue is stopped.
|
||||
if (entry.context.input_queue->IsStopped() &&
|
||||
entry.context.input_queue->Size() == 0) {
|
||||
for (auto& q : entry.context.output_queues) q->Stop();
|
||||
break;
|
||||
}
|
||||
}
|
||||
for (auto& q : entry.context.output_queues) q->Stop();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user