From 20f9a4ff2c30943edde67d583790cf166ec42fe5 Mon Sep 17 00:00:00 2001 From: sladro Date: Wed, 7 Jan 2026 18:46:53 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=BA=BA=E8=84=B8=E7=94=BB?= =?UTF-8?q?=E6=A1=86,xin1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- include/ai_scheduler.h | 5 ++ include/utils/simple_json.h | 4 + plugins/ai_face_det/ai_face_det_node.cpp | 84 +++++++++++++++++++- plugins/ai_face_recog/ai_face_recog_node.cpp | 79 +++++++++++++++++- plugins/preprocess/preprocess_node.cpp | 4 + src/ai_scheduler.cpp | 4 +- 6 files changed, 174 insertions(+), 6 deletions(-) diff --git a/include/ai_scheduler.h b/include/ai_scheduler.h index 8e0e8b5..b1a07ce 100644 --- a/include/ai_scheduler.h +++ b/include/ai_scheduler.h @@ -33,6 +33,11 @@ struct InferInput { int width = 0; int height = 0; bool is_nhwc = true; // true: NHWC, false: NCHW + +#if defined(RK3588_ENABLE_RKNN) + // The actual data type of `data` passed to RKNN. Default preserves existing behavior. + rknn_tensor_type type = RKNN_TENSOR_UINT8; +#endif }; struct InferOutput { diff --git a/include/utils/simple_json.h b/include/utils/simple_json.h index ad4843b..1235d21 100644 --- a/include/utils/simple_json.h +++ b/include/utils/simple_json.h @@ -4,12 +4,14 @@ #pragma once +#include #include #include #include #include #include #include +#include #include #include @@ -119,6 +121,8 @@ public: return child->AsString(def); } else if constexpr (std::is_same_v) { return child->AsNumber(def); + } else if constexpr (std::is_same_v) { + return static_cast(child->AsNumber(static_cast(def))); } else if constexpr (std::is_same_v) { return child->AsInt(def); } else if constexpr (std::is_same_v) { diff --git a/plugins/ai_face_det/ai_face_det_node.cpp b/plugins/ai_face_det/ai_face_det_node.cpp index f9fa5c3..144600c 100644 --- a/plugins/ai_face_det/ai_face_det_node.cpp +++ b/plugins/ai_face_det/ai_face_det_node.cpp @@ -327,6 +327,29 @@ public: input_format_ = fmt; for (auto& c : input_format_) c = static_cast(std::tolower(static_cast(c))); + input_dtype_ = config.ValueOr("input_dtype", input_dtype_); + for (auto& c : input_dtype_) c = static_cast(std::tolower(static_cast(c))); + + // Optional normalization when input_dtype is float. + // - scale/bias: x = x*scale + bias + // - mean/std: x = (x - mean[c]) / std[c] + if (const SimpleJson* norm = config.Find("normalize"); norm && norm->IsObject()) { + if (const SimpleJson* mean = norm->Find("mean"); mean && mean->IsArray() && mean->AsArray().size() >= 3) { + for (int i = 0; i < 3; ++i) { + norm_mean_[static_cast(i)] = static_cast(mean->AsArray()[static_cast(i)].AsNumber(norm_mean_[static_cast(i)])); + } + norm_use_mean_std_ = true; + } + if (const SimpleJson* st = norm->Find("std"); st && st->IsArray() && st->AsArray().size() >= 3) { + for (int i = 0; i < 3; ++i) { + norm_std_[static_cast(i)] = static_cast(st->AsArray()[static_cast(i)].AsNumber(norm_std_[static_cast(i)])); + } + norm_use_mean_std_ = true; + } + norm_scale_ = norm->ValueOr("scale", norm_scale_); + norm_bias_ = norm->ValueOr("bias", norm_bias_); + } + // RetinaFace priors defaults for 320 input (MobileNet0.25). steps_ = {8, 16, 32}; min_sizes_ = {{16, 32}, {64, 128}, {256, 512}}; @@ -407,6 +430,29 @@ public: max_faces_ = std::max(1, new_config.ValueOr("max_faces", max_faces_)); output_landmarks_ = new_config.ValueOr("output_landmarks", output_landmarks_); + std::string dtype = new_config.ValueOr("input_dtype", input_dtype_); + for (auto& c : dtype) c = static_cast(std::tolower(static_cast(c))); + input_dtype_ = std::move(dtype); + + if (const SimpleJson* norm = new_config.Find("normalize"); norm && norm->IsObject()) { + bool use_ms = false; + if (const SimpleJson* mean = norm->Find("mean"); mean && mean->IsArray() && mean->AsArray().size() >= 3) { + for (int i = 0; i < 3; ++i) { + norm_mean_[static_cast(i)] = static_cast(mean->AsArray()[static_cast(i)].AsNumber(norm_mean_[static_cast(i)])); + } + use_ms = true; + } + if (const SimpleJson* st = norm->Find("std"); st && st->IsArray() && st->AsArray().size() >= 3) { + for (int i = 0; i < 3; ++i) { + norm_std_[static_cast(i)] = static_cast(st->AsArray()[static_cast(i)].AsNumber(norm_std_[static_cast(i)])); + } + use_ms = true; + } + norm_use_mean_std_ = use_ms; + norm_scale_ = norm->ValueOr("scale", norm_scale_); + norm_bias_ = norm->ValueOr("bias", norm_bias_); + } + if (const SimpleJson* pri = new_config.Find("prior"); pri && pri->IsObject()) { if (const SimpleJson* steps = pri->Find("steps"); steps && steps->IsArray()) { std::vector new_steps; @@ -497,12 +543,37 @@ private: } InferInput input; - input.data = input_ptr; - input.size = in_size; input.width = in_w; input.height = in_h; input.is_nhwc = true; + // Default: keep existing UINT8 behavior. + if (input_dtype_ == "float" || input_dtype_ == "f32" || input_dtype_ == "float32") { + float_input_buf_.resize(static_cast(in_w) * static_cast(in_h) * 3); + const size_t pix = static_cast(in_w) * static_cast(in_h); + const uint8_t* p = reinterpret_cast(input_ptr); + for (size_t i = 0; i < pix; ++i) { + for (int c = 0; c < 3; ++c) { + float x = static_cast(p[i * 3 + static_cast(c)]); + if (norm_use_mean_std_) { + const float st = std::fabs(norm_std_[static_cast(c)]) < 1e-6f ? 1.0f : norm_std_[static_cast(c)]; + x = (x - norm_mean_[static_cast(c)]) / st; + } else { + x = x * norm_scale_ + norm_bias_; + } + float_input_buf_[i * 3 + static_cast(c)] = x; + } + } + + input.data = float_input_buf_.data(); + input.size = float_input_buf_.size() * sizeof(float); + input.type = RKNN_TENSOR_FLOAT32; + } else { + input.data = input_ptr; + input.size = in_size; + input.type = RKNN_TENSOR_UINT8; + } + auto r = AiScheduler::Instance().InferBorrowed(model_handle_, input); if (!r.success) { std::cerr << "[ai_face_det] inference failed: " << r.error << "\n"; @@ -699,6 +770,14 @@ private: std::string input_format_ = "rgb"; + // Model input dtype: "uint8" (default) or "float32". + std::string input_dtype_ = "uint8"; + float norm_scale_ = 1.0f; + float norm_bias_ = 0.0f; + bool norm_use_mean_std_ = false; + std::array norm_mean_{{0.0f, 0.0f, 0.0f}}; + std::array norm_std_{{1.0f, 1.0f, 1.0f}}; + std::vector steps_; std::vector> min_sizes_; @@ -706,6 +785,7 @@ private: std::vector>> output_queues_; std::vector input_buf_; + std::vector float_input_buf_; ModelHandle model_handle_ = kInvalidModelHandle; int model_w_ = 320; diff --git a/plugins/ai_face_recog/ai_face_recog_node.cpp b/plugins/ai_face_recog/ai_face_recog_node.cpp index f3d925e..f0c5b1e 100644 --- a/plugins/ai_face_recog/ai_face_recog_node.cpp +++ b/plugins/ai_face_recog/ai_face_recog_node.cpp @@ -569,6 +569,26 @@ public: model_input_format_ = fmt; for (auto& c : model_input_format_) c = static_cast(std::tolower(static_cast(c))); + input_dtype_ = config.ValueOr("input_dtype", input_dtype_); + for (auto& c : input_dtype_) c = static_cast(std::tolower(static_cast(c))); + + if (const SimpleJson* norm = config.Find("normalize"); norm && norm->IsObject()) { + if (const SimpleJson* mean = norm->Find("mean"); mean && mean->IsArray() && mean->AsArray().size() >= 3) { + for (int i = 0; i < 3; ++i) { + norm_mean_[static_cast(i)] = static_cast(mean->AsArray()[static_cast(i)].AsNumber(norm_mean_[static_cast(i)])); + } + norm_use_mean_std_ = true; + } + if (const SimpleJson* st = norm->Find("std"); st && st->IsArray() && st->AsArray().size() >= 3) { + for (int i = 0; i < 3; ++i) { + norm_std_[static_cast(i)] = static_cast(st->AsArray()[static_cast(i)].AsNumber(norm_std_[static_cast(i)])); + } + norm_use_mean_std_ = true; + } + norm_scale_ = norm->ValueOr("scale", norm_scale_); + norm_bias_ = norm->ValueOr("bias", norm_bias_); + } + // Gallery if (const SimpleJson* g = config.Find("gallery"); g && g->IsObject()) { gallery_backend_ = g->ValueOr("backend", gallery_backend_); @@ -643,6 +663,29 @@ public: thr_margin_ = th->ValueOr("margin", thr_margin_); } + std::string dtype = new_config.ValueOr("input_dtype", input_dtype_); + for (auto& c : dtype) c = static_cast(std::tolower(static_cast(c))); + input_dtype_ = std::move(dtype); + + if (const SimpleJson* norm = new_config.Find("normalize"); norm && norm->IsObject()) { + bool use_ms = false; + if (const SimpleJson* mean = norm->Find("mean"); mean && mean->IsArray() && mean->AsArray().size() >= 3) { + for (int i = 0; i < 3; ++i) { + norm_mean_[static_cast(i)] = static_cast(mean->AsArray()[static_cast(i)].AsNumber(norm_mean_[static_cast(i)])); + } + use_ms = true; + } + if (const SimpleJson* st = norm->Find("std"); st && st->IsArray() && st->AsArray().size() >= 3) { + for (int i = 0; i < 3; ++i) { + norm_std_[static_cast(i)] = static_cast(st->AsArray()[static_cast(i)].AsNumber(norm_std_[static_cast(i)])); + } + use_ms = true; + } + norm_use_mean_std_ = use_ms; + norm_scale_ = norm->ValueOr("scale", norm_scale_); + norm_bias_ = norm->ValueOr("bias", norm_bias_); + } + // Gallery updates bool reload = false; if (const SimpleJson* g = new_config.Find("gallery"); g && g->IsObject()) { @@ -773,12 +816,36 @@ private: } InferInput in; - in.data = face_buf_.data(); - in.size = face_buf_.size(); in.width = model_w_; in.height = model_h_; in.is_nhwc = true; + if (input_dtype_ == "float" || input_dtype_ == "f32" || input_dtype_ == "float32") { + float_input_buf_.resize(static_cast(model_w_) * static_cast(model_h_) * 3); + const size_t pix = static_cast(model_w_) * static_cast(model_h_); + const uint8_t* p = face_buf_.data(); + for (size_t ii = 0; ii < pix; ++ii) { + for (int c = 0; c < 3; ++c) { + float x = static_cast(p[ii * 3 + static_cast(c)]); + if (norm_use_mean_std_) { + const float st = std::fabs(norm_std_[static_cast(c)]) < 1e-6f ? 1.0f : norm_std_[static_cast(c)]; + x = (x - norm_mean_[static_cast(c)]) / st; + } else { + x = x * norm_scale_ + norm_bias_; + } + float_input_buf_[ii * 3 + static_cast(c)] = x; + } + } + + in.data = float_input_buf_.data(); + in.size = float_input_buf_.size() * sizeof(float); + in.type = RKNN_TENSOR_FLOAT32; + } else { + in.data = face_buf_.data(); + in.size = face_buf_.size(); + in.type = RKNN_TENSOR_UINT8; + } + auto r = AiScheduler::Instance().InferBorrowed(model_handle_, in); if (!r.success || r.outputs.empty()) { std::cerr << "[ai_face_recog] inference failed: " << (r.error.empty() ? "unknown" : r.error) << "\n"; @@ -831,6 +898,13 @@ private: std::string model_input_format_ = "rgb"; + std::string input_dtype_ = "uint8"; + float norm_scale_ = 1.0f; + float norm_bias_ = 0.0f; + bool norm_use_mean_std_ = false; + std::array norm_mean_{{0.0f, 0.0f, 0.0f}}; + std::array norm_std_{{1.0f, 1.0f, 1.0f}}; + std::string gallery_backend_ = "file"; std::string gallery_path_; // base path without extension: .json / .bin bool gallery_load_on_start_ = true; @@ -842,6 +916,7 @@ private: std::vector>> output_queues_; std::vector face_buf_; + std::vector float_input_buf_; ModelHandle model_handle_ = kInvalidModelHandle; int model_w_ = 112; diff --git a/plugins/preprocess/preprocess_node.cpp b/plugins/preprocess/preprocess_node.cpp index 604a21c..b60871b 100644 --- a/plugins/preprocess/preprocess_node.cpp +++ b/plugins/preprocess/preprocess_node.cpp @@ -574,6 +574,8 @@ private: out_frame->pts = frame->pts; out_frame->frame_id = frame->frame_id; out_frame->det = frame->det; + out_frame->face_det = frame->face_det; + out_frame->face_recog = frame->face_recog; out_frame->user_meta = frame->user_meta; SetupPlanes(*out_frame, out_fmt); @@ -658,6 +660,8 @@ private: out_frame->pts = frame->pts; out_frame->frame_id = frame->frame_id; out_frame->det = frame->det; + out_frame->face_det = frame->face_det; + out_frame->face_recog = frame->face_recog; out_frame->user_meta = frame->user_meta; SetupPlanes(*out_frame, out_fmt); diff --git a/src/ai_scheduler.cpp b/src/ai_scheduler.cpp index 89591aa..bdf0646 100644 --- a/src/ai_scheduler.cpp +++ b/src/ai_scheduler.cpp @@ -258,7 +258,7 @@ InferResult AiScheduler::Infer(ModelHandle handle, const InferInput& input) { rknn_input inputs[1]; memset(inputs, 0, sizeof(inputs)); inputs[0].index = 0; - inputs[0].type = RKNN_TENSOR_UINT8; + inputs[0].type = input.type; inputs[0].size = input.size; inputs[0].fmt = input.is_nhwc ? RKNN_TENSOR_NHWC : RKNN_TENSOR_NCHW; inputs[0].buf = const_cast(input.data); @@ -358,7 +358,7 @@ AiScheduler::BorrowedInferResult AiScheduler::InferBorrowed(ModelHandle handle, rknn_input inputs[1]; memset(inputs, 0, sizeof(inputs)); inputs[0].index = 0; - inputs[0].type = RKNN_TENSOR_UINT8; + inputs[0].type = input.type; inputs[0].size = input.size; inputs[0].fmt = input.is_nhwc ? RKNN_TENSOR_NHWC : RKNN_TENSOR_NCHW; inputs[0].buf = const_cast(input.data);