diff --git a/plugins/ai_shoe_det/ai_shoe_det_node.cpp b/plugins/ai_shoe_det/ai_shoe_det_node.cpp index d6cdc4e..fb67a1f 100644 --- a/plugins/ai_shoe_det/ai_shoe_det_node.cpp +++ b/plugins/ai_shoe_det/ai_shoe_det_node.cpp @@ -5,8 +5,10 @@ */ #include +#include #include #include +#include #include #include #include @@ -45,6 +47,226 @@ struct DynamicRoiConfig { float height_scale = 0.38f; }; +enum class V8BoxFormat { Auto, CxCyWh, XyXy, XyWh }; +enum class V8ClsActivation { Auto, None, Sigmoid }; + +uint32_t TensorTypeSizeBytes(rknn_tensor_type t) { + switch (t) { + case RKNN_TENSOR_INT8: + case RKNN_TENSOR_UINT8: + return 1; + case RKNN_TENSOR_FLOAT16: + return 2; + case RKNN_TENSOR_FLOAT32: + return 4; + default: + return 1; + } +} + +inline float Sigmoid(float x) { + return 1.0f / (1.0f + std::exp(-x)); +} + +inline float Fp16ToFp32(uint16_t h) { + const int sign = (h & 0x8000) ? -1 : 1; + const int exp = (h >> 10) & 0x1F; + const int mant = h & 0x03FF; + + if (exp == 0) { + if (mant == 0) return sign < 0 ? -0.0f : 0.0f; + return static_cast(sign) * std::ldexp(static_cast(mant), -24); + } + if (exp == 0x1F) { + if (mant == 0) return sign < 0 ? -INFINITY : INFINITY; + return std::numeric_limits::quiet_NaN(); + } + return static_cast(sign) * + std::ldexp(static_cast(mant + 1024), exp - 25); +} + +inline int32_t ClipFloat(float val, float min_val, float max_val) { + return static_cast(val <= min_val ? min_val : (val >= max_val ? max_val : val)); +} + +inline int8_t QuantizeF32ToAffine(float f32, int32_t zp, float scale) { + float dst_val = (f32 / scale) + zp; + return static_cast(ClipFloat(dst_val, -128, 127)); +} + +inline float DequantizeAffineToF32(int8_t qnt, int32_t zp, float scale) { + return (static_cast(qnt) - static_cast(zp)) * scale; +} + +int DefaultV8NumBoxes(int model_h, int model_w) { + if (model_h <= 0 || model_w <= 0) return 0; + return (model_h / 8) * (model_w / 8) + + (model_h / 16) * (model_w / 16) + + (model_h / 32) * (model_w / 32); +} + +struct V8LayoutInfo { + int num_boxes = 0; + bool channels_first = true; +}; + +float ScoreBoxCandidate(float x, float y, float w, float h, int model_w, int model_h) { + float s = 0.0f; + if (w > 0.0f && h > 0.0f) s += 3.0f; + if (w <= model_w * 1.2f) s += 1.0f; + if (h <= model_h * 1.2f) s += 1.0f; + if (x >= -model_w * 0.1f) s += 1.0f; + if (y >= -model_h * 0.1f) s += 1.0f; + if ((x + w) <= model_w * 1.2f) s += 1.0f; + if ((y + h) <= model_h * 1.2f) s += 1.0f; + return s; +} + +bool SeemsNormalized(float a, float b, float c, float d) { + auto in_range = [](float v) { return v >= -0.05f && v <= 2.5f; }; + return in_range(a) && in_range(b) && in_range(c) && in_range(d); +} + +void DecodeV8Box(float a, float b, float c, float d, int model_w, int model_h, V8BoxFormat fmt, + float& out_x, float& out_y, float& out_w, float& out_h) { + if (SeemsNormalized(a, b, c, d)) { + a *= static_cast(model_w); + b *= static_cast(model_h); + c *= static_cast(model_w); + d *= static_cast(model_h); + } + + auto decode_cxcywh = [&](float& x, float& y, float& w, float& h) { + x = a - c / 2.0f; + y = b - d / 2.0f; + w = c; + h = d; + }; + auto decode_xyxy = [&](float& x, float& y, float& w, float& h) { + x = a; + y = b; + w = c - a; + h = d - b; + }; + auto decode_xywh = [&](float& x, float& y, float& w, float& h) { + x = a; + y = b; + w = c; + h = d; + }; + + if (fmt == V8BoxFormat::CxCyWh) { + decode_cxcywh(out_x, out_y, out_w, out_h); + return; + } + if (fmt == V8BoxFormat::XyXy) { + decode_xyxy(out_x, out_y, out_w, out_h); + return; + } + if (fmt == V8BoxFormat::XyWh) { + decode_xywh(out_x, out_y, out_w, out_h); + return; + } + + float x1 = 0.0f, y1 = 0.0f, w1 = 0.0f, h1 = 0.0f; + float x2 = 0.0f, y2 = 0.0f, w2 = 0.0f, h2 = 0.0f; + float x3 = 0.0f, y3 = 0.0f, w3 = 0.0f, h3 = 0.0f; + decode_cxcywh(x1, y1, w1, h1); + decode_xyxy(x2, y2, w2, h2); + decode_xywh(x3, y3, w3, h3); + + const float s1 = ScoreBoxCandidate(x1, y1, w1, h1, model_w, model_h); + const float s2 = ScoreBoxCandidate(x2, y2, w2, h2, model_w, model_h); + const float s3 = ScoreBoxCandidate(x3, y3, w3, h3, model_w, model_h); + if (s2 >= s1 && s2 >= s3) { + out_x = x2; out_y = y2; out_w = w2; out_h = h2; + } else if (s3 >= s1 && s3 >= s2) { + out_x = x3; out_y = y3; out_w = w3; out_h = h3; + } else { + out_x = x1; out_y = y1; out_w = w1; out_h = h1; + } +} + +bool ResolveV8ApplySigmoid(const float* output, int num_boxes, int num_classes, bool channels_first, + V8ClsActivation act_mode) { + if (act_mode == V8ClsActivation::None) return false; + if (act_mode == V8ClsActivation::Sigmoid) return true; + if (!output || num_boxes <= 0 || num_classes <= 0) return false; + + const int num_channels = 4 + num_classes; + const int sample_boxes = std::min(num_boxes, 64); + float min_v = 1e9f; + float max_v = -1e9f; + for (int i = 0; i < sample_boxes; ++i) { + for (int c = 0; c < num_classes; ++c) { + const float v = channels_first ? output[(4 + c) * num_boxes + i] + : output[i * num_channels + (4 + c)]; + if (v < min_v) min_v = v; + if (v > max_v) max_v = v; + } + } + return (min_v < -0.1f || max_v > 1.5f); +} + +V8LayoutInfo ResolveV8Layout(const std::vector& dims, size_t byte_size, + rknn_tensor_type type, int num_classes, + int model_h, int model_w) { + V8LayoutInfo info; + const int num_channels = 4 + num_classes; + const uint32_t elem_bytes = TensorTypeSizeBytes(type); + const size_t total_elems = elem_bytes > 0 ? (byte_size / elem_bytes) : 0; + const size_t max_boxes_from_data = static_cast(num_channels) > 0 + ? (total_elems / static_cast(num_channels)) + : 0; + + int ch_idx = -1; + for (size_t i = 0; i < dims.size(); ++i) { + if (dims[i] == static_cast(num_channels)) { + ch_idx = static_cast(i); + break; + } + } + + if (ch_idx >= 0 && total_elems >= static_cast(num_channels)) { + info.num_boxes = static_cast(max_boxes_from_data); + int prev_non1 = 1; + for (int i = ch_idx - 1; i >= 0; --i) { + if (dims[static_cast(i)] > 1U) { + prev_non1 = static_cast(dims[static_cast(i)]); + break; + } + } + int next_non1 = 1; + for (size_t i = static_cast(ch_idx + 1); i < dims.size(); ++i) { + if (dims[i] > 1U) { + next_non1 = static_cast(dims[i]); + break; + } + } + if (next_non1 > 1 && prev_non1 <= 1) { + info.channels_first = true; + } else if (prev_non1 > 1 && next_non1 <= 1) { + info.channels_first = false; + } else if (next_non1 > 1 && prev_non1 > 1) { + info.channels_first = next_non1 >= prev_non1; + } + } + + if (info.num_boxes <= 0 && max_boxes_from_data > 0) { + info.num_boxes = static_cast(max_boxes_from_data); + } + if (info.num_boxes <= 0) { + info.num_boxes = DefaultV8NumBoxes(model_h, model_w); + } + if (info.num_boxes <= 0) { + info.num_boxes = 8400; + } + if (max_boxes_from_data > 0 && static_cast(info.num_boxes) > max_boxes_from_data) { + info.num_boxes = static_cast(max_boxes_from_data); + } + return info; +} + class AiShoeDetNode : public INode { public: std::string Id() const override { return id_; } @@ -322,30 +544,23 @@ private: int win_x, int win_y, int win_w, int win_h) { std::vector 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(outputs[0].data); - const size_t value_count = outputs[0].size / sizeof(float); - const int num_boxes = static_cast(std::min(8400, value_count / 5)); + constexpr int kNumClasses = 1; + const V8LayoutInfo layout = ResolveV8Layout( + outputs[0].dims, outputs[0].size, outputs[0].type, kNumClasses, model_h_, model_w_); + const int num_boxes = layout.num_boxes; if (num_boxes <= 0) return dets; float scale_x = static_cast(win_w) / model_w_; float scale_y = static_cast(win_h) / model_h_; + const int num_channels = 4 + kNumClasses; - 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; + auto append_box = [&](float a, float b, float c, float d, float conf) { + if (!std::isfinite(a) || !std::isfinite(b) || !std::isfinite(c) || !std::isfinite(d)) return; + float x = 0.0f, y = 0.0f, w = 0.0f, h = 0.0f; + DecodeV8Box(a, b, c, d, model_w_, model_h_, V8BoxFormat::Auto, x, y, w, h); + if (!std::isfinite(x) || !std::isfinite(y) || !std::isfinite(w) || !std::isfinite(h)) return; + if (w <= 1e-3f || h <= 1e-3f) return; + if (w > model_w_ * 1.2f || h > model_h_ * 1.2f) return; DetBox box; box.x = win_x + x * scale_x; @@ -355,6 +570,58 @@ private: box.conf = conf; box.class_id = dynamic_roi_.shoe_class_id; dets.push_back(box); + }; + + if (outputs[0].type == RKNN_TENSOR_FLOAT32) { + const float* data = reinterpret_cast(outputs[0].data); + const bool apply_sigmoid = ResolveV8ApplySigmoid( + data, num_boxes, kNumClasses, layout.channels_first, V8ClsActivation::Auto); + for (int i = 0; i < num_boxes; ++i) { + const float score = layout.channels_first ? data[4 * num_boxes + i] + : data[i * num_channels + 4]; + const float conf = apply_sigmoid ? Sigmoid(score) : score; + if (conf < conf_thresh_) continue; + const float a = layout.channels_first ? data[0 * num_boxes + i] : data[i * num_channels + 0]; + const float b = layout.channels_first ? data[1 * num_boxes + i] : data[i * num_channels + 1]; + const float c = layout.channels_first ? data[2 * num_boxes + i] : data[i * num_channels + 2]; + const float d = layout.channels_first ? data[3 * num_boxes + i] : data[i * num_channels + 3]; + append_box(a, b, c, d, conf); + } + } else if (outputs[0].type == RKNN_TENSOR_FLOAT16) { + const uint16_t* data = reinterpret_cast(outputs[0].data); + fp32_decode_buf_.resize(outputs[0].size / sizeof(uint16_t)); + for (size_t i = 0; i < fp32_decode_buf_.size(); ++i) { + fp32_decode_buf_[i] = Fp16ToFp32(data[i]); + } + const bool apply_sigmoid = ResolveV8ApplySigmoid( + fp32_decode_buf_.data(), num_boxes, kNumClasses, layout.channels_first, V8ClsActivation::Auto); + for (int i = 0; i < num_boxes; ++i) { + const float score = layout.channels_first ? fp32_decode_buf_[4 * num_boxes + i] + : fp32_decode_buf_[i * num_channels + 4]; + const float conf = apply_sigmoid ? Sigmoid(score) : score; + if (conf < conf_thresh_) continue; + const float a = layout.channels_first ? fp32_decode_buf_[0 * num_boxes + i] : fp32_decode_buf_[i * num_channels + 0]; + const float b = layout.channels_first ? fp32_decode_buf_[1 * num_boxes + i] : fp32_decode_buf_[i * num_channels + 1]; + const float c = layout.channels_first ? fp32_decode_buf_[2 * num_boxes + i] : fp32_decode_buf_[i * num_channels + 2]; + const float d = layout.channels_first ? fp32_decode_buf_[3 * num_boxes + i] : fp32_decode_buf_[i * num_channels + 3]; + append_box(a, b, c, d, conf); + } + } else if (outputs[0].type == RKNN_TENSOR_INT8) { + const int8_t* data = reinterpret_cast(outputs[0].data); + const int8_t thresh_i8 = QuantizeF32ToAffine(conf_thresh_, outputs[0].zp, outputs[0].scale); + for (int i = 0; i < num_boxes; ++i) { + const int8_t score_i8 = layout.channels_first ? data[4 * num_boxes + i] + : data[i * num_channels + 4]; + if (score_i8 < thresh_i8) continue; + const float conf = DequantizeAffineToF32(score_i8, outputs[0].zp, outputs[0].scale); + const float a = DequantizeAffineToF32(layout.channels_first ? data[0 * num_boxes + i] : data[i * num_channels + 0], outputs[0].zp, outputs[0].scale); + const float b = DequantizeAffineToF32(layout.channels_first ? data[1 * num_boxes + i] : data[i * num_channels + 1], outputs[0].zp, outputs[0].scale); + const float c = DequantizeAffineToF32(layout.channels_first ? data[2 * num_boxes + i] : data[i * num_channels + 2], outputs[0].zp, outputs[0].scale); + const float d = DequantizeAffineToF32(layout.channels_first ? data[3 * num_boxes + i] : data[i * num_channels + 3], outputs[0].zp, outputs[0].scale); + append_box(a, b, c, d, conf); + } + } else { + LogWarn("[ai_shoe_det] unsupported output tensor type"); } return dets; @@ -447,6 +714,7 @@ private: DynamicRoiConfig dynamic_roi_; std::vector windows_; std::vector input_buf_; + std::vector fp32_decode_buf_; std::shared_ptr> input_queue_; std::vector>> output_queues_;