Signed-off-by: Xing Song <xing_song@foxmail.com>
This commit is contained in:
Xing Song 2026-05-15 04:41:43 +08:00
parent 8b6decb418
commit 1cd7058905

View file

@ -156,7 +156,7 @@ cv::Mat quantize_input(const cv::Mat& float_img, float scale, int8_t zero_point)
}
static std::vector<Detection> get_detections(float* output, std::tuple<int, int, int> output_shape,
int stride, float conf_thresh) {
int stride, float conf_thresh, int index) {
std::vector<Detection> detections;
int grid_h = std::get<0>(output_shape);
@ -208,6 +208,8 @@ static std::vector<Detection> get_detections(float* output, std::tuple<int, int,
}
}
float keypoint_index = index + (i * grid_w) + j;
// Convert to absolute coordinates
float anchor_x = (j + 0.5f) * stride;
float anchor_y = (i + 0.5f) * stride;
@ -222,29 +224,33 @@ static std::vector<Detection> get_detections(float* output, std::tuple<int, int,
float x2 = anchor_x + right * stride;
float y2 = anchor_y + bottom * stride;
detections.push_back({x1, y1, x2, y2, max_score, class_id});
detections.push_back({x1, y1, x2, y2, max_score, class_id, keypoint_index});
}
}
index += grid_h * grid_w;
return detections;
}
std::vector<Detection> postprocess(std::tuple<float*, std::tuple<int, int, int>, int> out0,
std::tuple<float*, std::tuple<int, int, int>, int> out1,
std::tuple<float*, std::tuple<int, int, int>, int> out2,
std::tuple<float*, std::tuple<int, int, int>, int> out3,
std::tuple<cv::Mat, float, std::tuple<int, int>> input_tuple,
float conf_thresh, float iou_threshold) {
float scale = std::get<1>(input_tuple);
int pad_left = std::get<0>(std::get<2>(input_tuple));
int pad_top = std::get<1>(std::get<2>(input_tuple));
int index = 0;
std::vector<Detection> detections;
auto process_out = [&](auto& out) {
auto process_out = [&](auto& out, auto& index) {
float* output = std::get<0>(out);
auto shape = std::get<1>(out);
int stride = std::get<2>(out);
std::vector<Detection> dets = get_detections(output, shape, stride, conf_thresh);
std::vector<Detection> dets = get_detections(output, shape, stride, conf_thresh, index);
detections.insert(detections.end(), dets.begin(), dets.end());
index +=
};
// Process all three scales
@ -255,6 +261,13 @@ std::vector<Detection> postprocess(std::tuple<float*, std::tuple<int, int, int>,
// Map coordinates back to original image
std::vector<Detection> detections_orig;
for (const auto& det : detections) {
for (int j = 0; j < 17; i++) {
Keypoint* keypoint = det.keypoint;
keypoint[j].x = (out3[j*3*8400+0*8400+det.keypoint_index] - pad_left) / scale;
keypoint[j].y = (out3[j*3*8400+1*8400+det.keypoint_index] - pad_left) / scale;
keypoint[j].score = out3[j*3*8400+2*8400+det.keypoint_index];
}
float x1_orig = (det.x1 - pad_left) / scale;
float y1_orig = (det.y1 - pad_top) / scale;
float x2_orig = (det.x2 - pad_left) / scale;
@ -266,7 +279,7 @@ std::vector<Detection> postprocess(std::tuple<float*, std::tuple<int, int, int>,
x2_orig = std::max(0.0f, x2_orig);
y2_orig = std::max(0.0f, y2_orig);
detections_orig.push_back({x1_orig, y1_orig, x2_orig, y2_orig, det.score, det.class_id});
detections_orig.push_back({x1_orig, y1_orig, x2_orig, y2_orig, det.score, det.class_id, det.keypoint_index, det.keypoint});
}
// Apply NMS