Doni/safesight-edge
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve YOLOv8 parsing in shoe detector

Browse Source
This commit is contained in:
tian 2026-03-14 13:41:24 +08:00
parent 3383e13e88
commit f71c7eee74
1 changed files with 287 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

306
plugins/ai_shoe_det/ai_shoe_det_node.cpp
View File

@ -5,8 +5,10 @@
*/
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <limits>
#include <memory>
#include <string>
#include <vector>
@ -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<float>(sign) * std::ldexp(static_cast<float>(mant), -24);
}
if (exp == 0x1F) {
if (mant == 0) return sign < 0 ? -INFINITY : INFINITY;
return std::numeric_limits<float>::quiet_NaN();
}
return static_cast<float>(sign) *
std::ldexp(static_cast<float>(mant + 1024), exp - 25);
}
inline int32_t ClipFloat(float val, float min_val, float max_val) {
return static_cast<int32_t>(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<int8_t>(ClipFloat(dst_val, -128, 127));
}
inline float DequantizeAffineToF32(int8_t qnt, int32_t zp, float scale) {
return (static_cast<float>(qnt) - static_cast<float>(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<float>(model_w);
b *= static_cast<float>(model_h);
c *= static_cast<float>(model_w);
d *= static_cast<float>(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<uint32_t>& 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<size_t>(num_channels) > 0
? (total_elems / static_cast<size_t>(num_channels))
: 0;
int ch_idx = -1;
for (size_t i = 0; i < dims.size(); ++i) {
if (dims[i] == static_cast<uint32_t>(num_channels)) {
ch_idx = static_cast<int>(i);
break;
}
}
if (ch_idx >= 0 && total_elems >= static_cast<size_t>(num_channels)) {
info.num_boxes = static_cast<int>(max_boxes_from_data);
int prev_non1 = 1;
for (int i = ch_idx - 1; i >= 0; --i) {
if (dims[static_cast<size_t>(i)] > 1U) {
prev_non1 = static_cast<int>(dims[static_cast<size_t>(i)]);
break;
}
}
int next_non1 = 1;
for (size_t i = static_cast<size_t>(ch_idx + 1); i < dims.size(); ++i) {
if (dims[i] > 1U) {
next_non1 = static_cast<int>(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<int>(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<size_t>(info.num_boxes) > max_boxes_from_data) {
info.num_boxes = static_cast<int>(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<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));
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<float>(win_w) / model_w_;
float scale_y = static_cast<float>(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<const float*>(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<const uint16_t*>(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<const int8_t*>(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<DetWindow> windows_;
std::vector<uint8_t> input_buf_;
std::vector<float> fp32_decode_buf_;
std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;