712 lines
40 KiB
C++
712 lines
40 KiB
C++
#include "yolo_runtime_backend.h"
|
|
#include "yolo_cpu_image_utils.h"
|
|
#include "yolo_letterbox.h"
|
|
#include "yolo_nms_iou.h"
|
|
#include "yolo_onnx_input_utils.h"
|
|
|
|
#include <onnxruntime_cxx_api.h>
|
|
|
|
#include <algorithm>
|
|
#include <cctype>
|
|
#include <chrono>
|
|
#include <cstdio>
|
|
#include <cmath>
|
|
#include <cstring>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <numeric>
|
|
#include <sstream>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <thread>
|
|
#include <vector>
|
|
|
|
namespace fs = std::filesystem;
|
|
|
|
namespace {
|
|
|
|
struct runtime_config {
|
|
int width = 640;
|
|
int height = 640;
|
|
bool bgr = false;
|
|
std::string layout = "nchw";
|
|
std::vector<float> mean{0.0f, 0.0f, 0.0f};
|
|
std::vector<float> standard_deviation{255.0f, 255.0f, 255.0f};
|
|
std::vector<std::string> labels;
|
|
yolo_task_type_t task = YOLO_TASK_DETECT;
|
|
};
|
|
|
|
struct candidate {
|
|
int class_id = 0;
|
|
float score = 0.0f;
|
|
float x = 0.0f;
|
|
float y = 0.0f;
|
|
float width = 0.0f;
|
|
float height = 0.0f;
|
|
float angle = 0.0f;
|
|
std::vector<float> mask_coefficients;
|
|
std::vector<yolo_keypoint_v1_t> keypoints;
|
|
};
|
|
|
|
std::string trim(std::string value) {
|
|
const auto first = value.find_first_not_of(" \t\r\n\"'");
|
|
if (first == std::string::npos) return {};
|
|
const auto last = value.find_last_not_of(" \t\r\n\"'");
|
|
return value.substr(first, last - first + 1);
|
|
}
|
|
|
|
std::map<std::string, std::string> read_yaml(const fs::path& path) {
|
|
std::ifstream input(path);
|
|
if (!input) throw std::runtime_error("cannot open model config: " + path.string());
|
|
std::map<std::string, std::string> values;
|
|
std::map<int, std::string> sections;
|
|
std::string line;
|
|
while (std::getline(input, line)) {
|
|
const auto comment = line.find('#');
|
|
if (comment != std::string::npos) line.erase(comment);
|
|
if (trim(line).empty()) continue;
|
|
const int indent = static_cast<int>(line.find_first_not_of(" \t"));
|
|
const std::string content = trim(line);
|
|
const auto separator = content.find(':');
|
|
if (separator == std::string::npos) continue;
|
|
for (auto it = sections.begin(); it != sections.end();) {
|
|
if (it->first >= indent) it = sections.erase(it); else ++it;
|
|
}
|
|
const std::string key = trim(content.substr(0, separator));
|
|
const std::string value = trim(content.substr(separator + 1));
|
|
std::string full_key;
|
|
for (const auto& section : sections) {
|
|
if (!full_key.empty()) full_key += '.';
|
|
full_key += section.second;
|
|
}
|
|
if (!full_key.empty()) full_key += '.';
|
|
full_key += key;
|
|
if (value.empty()) sections[indent] = key;
|
|
else values[full_key] = value;
|
|
}
|
|
return values;
|
|
}
|
|
|
|
std::vector<std::string> split_list(std::string value) {
|
|
value = trim(value);
|
|
if (!value.empty() && value.front() == '[') value.erase(value.begin());
|
|
if (!value.empty() && value.back() == ']') value.pop_back();
|
|
std::vector<std::string> items;
|
|
std::stringstream stream(value);
|
|
std::string item;
|
|
while (std::getline(stream, item, ',')) {
|
|
item = trim(item);
|
|
if (!item.empty()) items.push_back(item);
|
|
}
|
|
return items;
|
|
}
|
|
|
|
std::vector<float> float_list(const std::map<std::string, std::string>& values, const char* key, std::vector<float> fallback) {
|
|
const auto found = values.find(key);
|
|
if (found == values.end()) return fallback;
|
|
std::vector<float> result;
|
|
for (const auto& item : split_list(found->second)) {
|
|
try { result.push_back(std::stof(item)); } catch (...) { return fallback; }
|
|
}
|
|
return result.empty() ? fallback : result;
|
|
}
|
|
|
|
runtime_config load_config(const fs::path& config_path) {
|
|
const auto values = read_yaml(config_path);
|
|
runtime_config config;
|
|
if (const auto it = values.find("task"); it != values.end()) {
|
|
auto task = trim(it->second);
|
|
std::transform(task.begin(), task.end(), task.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
|
if (task == "cls" || task == "classify" || task == "classification") config.task = YOLO_TASK_CLS;
|
|
else if (task == "obb") config.task = YOLO_TASK_OBB;
|
|
else if (task == "seg" || task == "segment") config.task = YOLO_TASK_SEG;
|
|
else if (task == "pose") config.task = YOLO_TASK_POSE;
|
|
else if (task != "detect") throw std::runtime_error("Windows CPU Runtime supports Detect, CLS, OBB, Seg, and Pose models only");
|
|
}
|
|
if (const auto it = values.find("input.width"); it != values.end()) config.width = std::stoi(it->second);
|
|
if (const auto it = values.find("input.height"); it != values.end()) config.height = std::stoi(it->second);
|
|
if (const auto it = values.find("input.format"); it != values.end()) {
|
|
auto format = trim(it->second);
|
|
std::transform(format.begin(), format.end(), format.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
|
config.bgr = format == "bgr";
|
|
}
|
|
if (const auto it = values.find("input.layout"); it != values.end()) config.layout = trim(it->second);
|
|
config.mean = float_list(values, "input.normalization.mean", config.mean);
|
|
config.standard_deviation = float_list(values, "input.normalization.std", config.standard_deviation);
|
|
if (const auto it = values.find("labels"); it != values.end()) config.labels = split_list(it->second);
|
|
if (const auto it = values.find("labels_file"); it != values.end()) {
|
|
fs::path labels_path(trim(it->second));
|
|
if (labels_path.is_relative()) labels_path = config_path.parent_path() / labels_path;
|
|
std::ifstream labels(labels_path);
|
|
std::string label;
|
|
while (std::getline(labels, label)) {
|
|
label = trim(label);
|
|
if (!label.empty()) config.labels.push_back(label);
|
|
}
|
|
}
|
|
if (config.width <= 0 || config.height <= 0) throw std::runtime_error("invalid configured model dimensions");
|
|
return config;
|
|
}
|
|
|
|
float bilinear_channel(
|
|
const yolo_image_view_v1_t& image,
|
|
float source_x,
|
|
float source_y,
|
|
int channel,
|
|
bool quantize_to_u8) {
|
|
const int x0 = std::clamp(static_cast<int>(std::floor(source_x)), 0, image.width - 1);
|
|
const int y0 = std::clamp(static_cast<int>(std::floor(source_y)), 0, image.height - 1);
|
|
const int x1 = std::min(x0 + 1, image.width - 1);
|
|
const int y1 = std::min(y0 + 1, image.height - 1);
|
|
const float fx = source_x - x0;
|
|
const float fy = source_y - y0;
|
|
float r, g, b;
|
|
yolo_cpu_image::pixel_rgb(image, x0, y0, r, g, b); const float top_left = channel == 0 ? r : channel == 1 ? g : b;
|
|
yolo_cpu_image::pixel_rgb(image, x1, y0, r, g, b); const float top_right = channel == 0 ? r : channel == 1 ? g : b;
|
|
yolo_cpu_image::pixel_rgb(image, x0, y1, r, g, b); const float bottom_left = channel == 0 ? r : channel == 1 ? g : b;
|
|
yolo_cpu_image::pixel_rgb(image, x1, y1, r, g, b); const float bottom_right = channel == 0 ? r : channel == 1 ? g : b;
|
|
if (quantize_to_u8) {
|
|
return yolo_letterbox::bilinear_u8(
|
|
static_cast<unsigned char>(top_left),
|
|
static_cast<unsigned char>(top_right),
|
|
static_cast<unsigned char>(bottom_left),
|
|
static_cast<unsigned char>(bottom_right),
|
|
fx,
|
|
fy);
|
|
}
|
|
const float top = top_left * (1.0f - fx) + top_right * fx;
|
|
const float bottom = bottom_left * (1.0f - fx) + bottom_right * fx;
|
|
return top * (1.0f - fy) + bottom * fy;
|
|
}
|
|
|
|
yolo_nms::box nms_box(const candidate& value) {
|
|
return {value.x, value.y, value.width, value.height, value.angle};
|
|
}
|
|
|
|
} // namespace
|
|
|
|
struct yolo_runtime_backend {
|
|
std::unique_ptr<Ort::Env> environment;
|
|
Ort::SessionOptions session_options;
|
|
std::unique_ptr<Ort::Session> session;
|
|
Ort::MemoryInfo memory_info{nullptr};
|
|
std::string input_name;
|
|
std::vector<std::string> output_names;
|
|
std::vector<int64_t> input_shape;
|
|
runtime_config config;
|
|
float confidence_threshold = 0.25f;
|
|
float nms_threshold = 0.45f;
|
|
bool class_agnostic_nms = false;
|
|
bool nchw = true;
|
|
std::string last_error;
|
|
std::mutex mutex;
|
|
std::vector<std::vector<unsigned char>> segmentation_masks;
|
|
};
|
|
|
|
extern "C" int yolo_backend_get_abi_version(void) { return YOLO_RUNTIME_BACKEND_ABI_VERSION; }
|
|
|
|
extern "C" int yolo_backend_create(
|
|
const char* model_path,
|
|
const char* config_path,
|
|
float confidence_threshold,
|
|
float nms_threshold,
|
|
int class_agnostic_nms,
|
|
yolo_runtime_backend_t** output) {
|
|
if (!model_path || !config_path || !output) return YOLO_STATUS_INVALID_ARGUMENT;
|
|
try {
|
|
auto backend = std::make_unique<yolo_runtime_backend>();
|
|
backend->config = load_config(fs::u8path(config_path));
|
|
if (backend->config.task != YOLO_TASK_DETECT && backend->config.task != YOLO_TASK_CLS &&
|
|
backend->config.task != YOLO_TASK_OBB && backend->config.task != YOLO_TASK_SEG &&
|
|
backend->config.task != YOLO_TASK_POSE)
|
|
throw std::runtime_error("Windows CPU Runtime supports Detect, CLS, OBB, Seg, and Pose models only");
|
|
backend->confidence_threshold = confidence_threshold;
|
|
backend->nms_threshold = nms_threshold;
|
|
backend->class_agnostic_nms = class_agnostic_nms != 0;
|
|
backend->environment = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "rk3588_ai_windows_cpu");
|
|
backend->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
|
backend->session_options.SetIntraOpNumThreads(std::max(1u, std::thread::hardware_concurrency() / 2));
|
|
backend->session = std::make_unique<Ort::Session>(*backend->environment, fs::u8path(model_path).c_str(), backend->session_options);
|
|
if (backend->session->GetInputCount() != 1) throw std::runtime_error("Windows CPU Runtime requires exactly one ONNX input");
|
|
Ort::AllocatorWithDefaultOptions allocator;
|
|
auto input_name = backend->session->GetInputNameAllocated(0, allocator);
|
|
backend->input_name = input_name.get();
|
|
backend->input_shape = backend->session->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
|
|
if (backend->input_shape.size() != 4) throw std::runtime_error("Windows CPU Runtime requires a four-dimensional image input");
|
|
backend->nchw = yolo_onnx_input::resolve_nchw(backend->input_shape, backend->config.layout);
|
|
yolo_onnx_input::validate_three_channels(backend->input_shape, backend->nchw);
|
|
const size_t outputs = backend->session->GetOutputCount();
|
|
if (outputs == 0) throw std::runtime_error("ONNX model has no outputs");
|
|
for (size_t index = 0; index < outputs; ++index) {
|
|
auto output_name = backend->session->GetOutputNameAllocated(index, allocator);
|
|
backend->output_names.emplace_back(output_name.get());
|
|
}
|
|
backend->memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
|
*output = backend.release();
|
|
return YOLO_STATUS_OK;
|
|
} catch (const std::exception& exception) {
|
|
return YOLO_STATUS_INITIALIZATION_FAILED;
|
|
}
|
|
}
|
|
|
|
extern "C" int yolo_backend_infer(yolo_runtime_backend_t* backend, const yolo_image_view_v1_t* image, yolo_backend_result_list_t* results) {
|
|
if (!backend || !yolo_cpu_image::valid(image) || !results) return YOLO_STATUS_INVALID_ARGUMENT;
|
|
try {
|
|
std::lock_guard<std::mutex> lock(backend->mutex);
|
|
std::memset(results, 0, sizeof(*results));
|
|
results->task = backend->config.task;
|
|
const auto total_start = std::chrono::steady_clock::now();
|
|
|
|
const bool nchw = backend->nchw;
|
|
const int model_height = nchw && backend->input_shape[2] > 0 ? static_cast<int>(backend->input_shape[2])
|
|
: !nchw && backend->input_shape[1] > 0 ? static_cast<int>(backend->input_shape[1]) : backend->config.height;
|
|
const int model_width = nchw && backend->input_shape[3] > 0 ? static_cast<int>(backend->input_shape[3])
|
|
: !nchw && backend->input_shape[2] > 0 ? static_cast<int>(backend->input_shape[2]) : backend->config.width;
|
|
if (model_width <= 0 || model_height <= 0) throw std::runtime_error("invalid ONNX input dimensions");
|
|
const auto preprocess_start = std::chrono::steady_clock::now();
|
|
const bool classification = backend->config.task == YOLO_TASK_CLS;
|
|
const auto letterbox = classification
|
|
? yolo_letterbox::plan{1.0f, model_width, model_height, 0, 0}
|
|
: yolo_letterbox::make_plan(image->width, image->height, model_width, model_height);
|
|
const float scale = letterbox.scale;
|
|
const int resized_width = letterbox.resized_width;
|
|
const int resized_height = letterbox.resized_height;
|
|
const int pad_x = letterbox.pad_x;
|
|
const int pad_y = letterbox.pad_y;
|
|
const size_t values = static_cast<size_t>(model_width) * model_height * 3;
|
|
std::vector<float> tensor(values, 0.0f);
|
|
auto write = [&](int channel, int x, int y, float value) {
|
|
const float mean = backend->config.mean.size() > static_cast<size_t>(channel) ? backend->config.mean[channel] : 0.0f;
|
|
const float standard_deviation = backend->config.standard_deviation.size() > static_cast<size_t>(channel) && std::abs(backend->config.standard_deviation[channel]) > 0.000001f
|
|
? backend->config.standard_deviation[channel] : 1.0f;
|
|
const size_t index = nchw ? (static_cast<size_t>(channel) * model_height + y) * model_width + x
|
|
: (static_cast<size_t>(y) * model_width + x) * 3 + channel;
|
|
tensor[index] = (value - mean) / standard_deviation;
|
|
};
|
|
for (int y = 0; y < model_height; ++y)
|
|
for (int x = 0; x < model_width; ++x)
|
|
for (int channel = 0; channel < 3; ++channel) write(channel, x, y, 114.0f);
|
|
for (int y = 0; y < resized_height; ++y) {
|
|
const float source_y = yolo_letterbox::source_coordinate(y, image->height, resized_height);
|
|
for (int x = 0; x < resized_width; ++x) {
|
|
const float source_x = yolo_letterbox::source_coordinate(x, image->width, resized_width);
|
|
const float red = bilinear_channel(*image, source_x, source_y, 0, !classification);
|
|
const float green = bilinear_channel(*image, source_x, source_y, 1, !classification);
|
|
const float blue = bilinear_channel(*image, source_x, source_y, 2, !classification);
|
|
write(0, x + pad_x, y + pad_y, backend->config.bgr ? blue : red);
|
|
write(1, x + pad_x, y + pad_y, green);
|
|
write(2, x + pad_x, y + pad_y, backend->config.bgr ? red : blue);
|
|
}
|
|
}
|
|
results->performance.preprocess_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - preprocess_start).count();
|
|
|
|
std::vector<int64_t> input_shape = backend->input_shape;
|
|
for (auto& dimension : input_shape) if (dimension <= 0) dimension = 1;
|
|
if (nchw) { input_shape[1] = 3; input_shape[2] = model_height; input_shape[3] = model_width; }
|
|
else { input_shape[1] = model_height; input_shape[2] = model_width; input_shape[3] = 3; }
|
|
auto input_tensor = Ort::Value::CreateTensor<float>(backend->memory_info, tensor.data(), tensor.size(), input_shape.data(), input_shape.size());
|
|
std::vector<const char*> output_names;
|
|
for (const auto& name : backend->output_names) output_names.push_back(name.c_str());
|
|
const char* input_names[] = { backend->input_name.c_str() };
|
|
const auto inference_start = std::chrono::steady_clock::now();
|
|
auto output_values = backend->session->Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names.data(), output_names.size());
|
|
results->performance.inference_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - inference_start).count();
|
|
|
|
const auto postprocess_start = std::chrono::steady_clock::now();
|
|
if (classification) {
|
|
Ort::Value* classification_value = nullptr;
|
|
for (auto& output : output_values) {
|
|
if (!output.IsTensor()) continue;
|
|
const auto info = output.GetTensorTypeAndShapeInfo();
|
|
if (info.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT && info.GetElementCount() > 0) {
|
|
classification_value = &output;
|
|
break;
|
|
}
|
|
}
|
|
if (!classification_value) throw std::runtime_error("cannot find FP32 classification output");
|
|
const auto output_info = classification_value->GetTensorTypeAndShapeInfo();
|
|
const size_t class_count = output_info.GetElementCount();
|
|
const float* raw_scores = classification_value->GetTensorData<float>();
|
|
std::vector<float> scores(raw_scores, raw_scores + class_count);
|
|
for (const float score : scores)
|
|
if (!std::isfinite(score)) throw std::runtime_error("classification output contains a non-finite score");
|
|
|
|
const double raw_sum = std::accumulate(scores.begin(), scores.end(), 0.0);
|
|
const bool probabilities = std::all_of(scores.begin(), scores.end(), [](float score) {
|
|
return score >= 0.0f && score <= 1.0f;
|
|
}) && std::abs(raw_sum - 1.0) < 0.01;
|
|
if (!probabilities) {
|
|
const float maximum = *std::max_element(scores.begin(), scores.end());
|
|
double total = 0.0;
|
|
for (float& score : scores) {
|
|
score = std::exp(score - maximum);
|
|
total += score;
|
|
}
|
|
if (total <= 0.0) throw std::runtime_error("classification Softmax total is zero");
|
|
for (float& score : scores) score = static_cast<float>(score / total);
|
|
}
|
|
|
|
std::vector<size_t> order(class_count);
|
|
std::iota(order.begin(), order.end(), 0);
|
|
std::sort(order.begin(), order.end(), [&](size_t first, size_t second) {
|
|
return scores[first] > scores[second];
|
|
});
|
|
const size_t result_count = std::min<size_t>({5, class_count, YOLO_BACKEND_MAX_RESULTS});
|
|
for (size_t index = 0; index < result_count; ++index) {
|
|
const size_t class_id = order[index];
|
|
auto& result = results->results[results->count++];
|
|
result.task = YOLO_TASK_CLS;
|
|
result.class_id = static_cast<int>(class_id);
|
|
result.score = scores[class_id];
|
|
const std::string label = class_id < backend->config.labels.size()
|
|
? backend->config.labels[class_id] : "class_" + std::to_string(class_id);
|
|
std::snprintf(result.label, sizeof(result.label), "%s", label.c_str());
|
|
}
|
|
results->performance.postprocess_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - postprocess_start).count();
|
|
results->performance.total_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - total_start).count();
|
|
return YOLO_STATUS_OK;
|
|
}
|
|
|
|
if (backend->config.task == YOLO_TASK_POSE) {
|
|
Ort::Value* pose_keypoints = nullptr;
|
|
std::vector<Ort::Value*> pose_heads;
|
|
for (auto& output_value : output_values) {
|
|
if (!output_value.IsTensor()) continue;
|
|
const auto info = output_value.GetTensorTypeAndShapeInfo();
|
|
const auto shape = info.GetShape();
|
|
if (shape.size() != 4 || info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) continue;
|
|
if (shape[2] == 3 && shape[1] > 0 && shape[1] <= YOLO_BACKEND_MAX_KEYPOINTS)
|
|
pose_keypoints = &output_value;
|
|
else if (shape[1] > 64)
|
|
pose_heads.push_back(&output_value);
|
|
}
|
|
if (pose_keypoints && !pose_heads.empty()) {
|
|
const auto keypoint_shape = pose_keypoints->GetTensorTypeAndShapeInfo().GetShape();
|
|
const int pose_keypoint_count = static_cast<int>(keypoint_shape[1]);
|
|
const int total_positions = static_cast<int>(keypoint_shape[3]);
|
|
const float* keypoint_data = pose_keypoints->GetTensorData<float>();
|
|
std::sort(pose_heads.begin(), pose_heads.end(), [](Ort::Value* first, Ort::Value* second) {
|
|
return first->GetTensorTypeAndShapeInfo().GetShape()[2] >
|
|
second->GetTensorTypeAndShapeInfo().GetShape()[2];
|
|
});
|
|
std::vector<candidate> pose_candidates;
|
|
int position_offset = 0;
|
|
for (Ort::Value* head : pose_heads) {
|
|
const auto shape = head->GetTensorTypeAndShapeInfo().GetShape();
|
|
const int head_channels = static_cast<int>(shape[1]);
|
|
const int grid_height = static_cast<int>(shape[2]);
|
|
const int grid_width = static_cast<int>(shape[3]);
|
|
const int head_positions = grid_height * grid_width;
|
|
const int pose_class_count = head_channels - 64;
|
|
if (pose_class_count <= 0 || position_offset + head_positions > total_positions)
|
|
throw std::runtime_error("invalid split pose output shape");
|
|
const int stride = backend->config.height / grid_height;
|
|
const float* head_data = head->GetTensorData<float>();
|
|
const auto head_at = [&](int channel, int position) -> float {
|
|
return head_data[static_cast<size_t>(channel) * head_positions + position];
|
|
};
|
|
for (int position = 0; position < head_positions; ++position) {
|
|
int class_id = 0;
|
|
float score = 1.0f / (1.0f + std::exp(-head_at(64, position)));
|
|
for (int class_index = 1; class_index < pose_class_count; ++class_index) {
|
|
const float value = 1.0f / (1.0f + std::exp(-head_at(64 + class_index, position)));
|
|
if (value > score) { score = value; class_id = class_index; }
|
|
}
|
|
if (score < backend->confidence_threshold) continue;
|
|
float distances[4]{};
|
|
for (int side = 0; side < 4; ++side) {
|
|
float maximum = head_at(side * 16, position);
|
|
for (int bin = 1; bin < 16; ++bin)
|
|
maximum = std::max(maximum, head_at(side * 16 + bin, position));
|
|
float denominator = 0.0f;
|
|
float numerator = 0.0f;
|
|
for (int bin = 0; bin < 16; ++bin) {
|
|
const float probability = std::exp(head_at(side * 16 + bin, position) - maximum);
|
|
denominator += probability;
|
|
numerator += probability * bin;
|
|
}
|
|
distances[side] = numerator / denominator;
|
|
}
|
|
const int grid_x = position % grid_width;
|
|
const int grid_y = position / grid_width;
|
|
const float left = (grid_x + 0.5f - distances[0]) * stride;
|
|
const float top = (grid_y + 0.5f - distances[1]) * stride;
|
|
candidate item{
|
|
class_id, score,
|
|
std::clamp((left - pad_x) / scale, 0.0f, static_cast<float>(image->width)),
|
|
std::clamp((top - pad_y) / scale, 0.0f, static_cast<float>(image->height)),
|
|
std::clamp((distances[0] + distances[2]) * stride / scale, 0.0f, static_cast<float>(image->width)),
|
|
std::clamp((distances[1] + distances[3]) * stride / scale, 0.0f, static_cast<float>(image->height))
|
|
};
|
|
item.keypoints.resize(static_cast<size_t>(pose_keypoint_count));
|
|
const int global_position = position_offset + position;
|
|
for (int keypoint = 0; keypoint < pose_keypoint_count; ++keypoint) {
|
|
const size_t base = static_cast<size_t>(keypoint) * 3 * total_positions;
|
|
item.keypoints[static_cast<size_t>(keypoint)] = {
|
|
std::clamp((keypoint_data[base + global_position] - pad_x) / scale, 0.0f, static_cast<float>(image->width)),
|
|
std::clamp((keypoint_data[base + total_positions + global_position] - pad_y) / scale, 0.0f, static_cast<float>(image->height)),
|
|
keypoint_data[base + 2 * total_positions + global_position]
|
|
};
|
|
}
|
|
pose_candidates.push_back(std::move(item));
|
|
}
|
|
position_offset += head_positions;
|
|
}
|
|
std::sort(pose_candidates.begin(), pose_candidates.end(), [](const candidate& first, const candidate& second) {
|
|
return first.score > second.score;
|
|
});
|
|
std::vector<candidate> selected_pose;
|
|
for (const auto& item : pose_candidates) {
|
|
if (selected_pose.size() >= YOLO_BACKEND_MAX_RESULTS) break;
|
|
const bool suppressed = std::any_of(selected_pose.begin(), selected_pose.end(), [&](const candidate& existing) {
|
|
return (backend->class_agnostic_nms || existing.class_id == item.class_id) &&
|
|
yolo_nms::overlap_iou(
|
|
backend->config.task, nms_box(existing), nms_box(item)) > backend->nms_threshold;
|
|
});
|
|
if (!suppressed) selected_pose.push_back(item);
|
|
}
|
|
for (const auto& item : selected_pose) {
|
|
auto& result = results->results[results->count++];
|
|
result.task = YOLO_TASK_POSE;
|
|
result.class_id = item.class_id;
|
|
result.score = item.score;
|
|
result.bbox = {item.x, item.y, item.width, item.height, 0.0f, 0};
|
|
result.keypoint_count = static_cast<int>(item.keypoints.size());
|
|
std::copy(item.keypoints.begin(), item.keypoints.end(), result.keypoints);
|
|
const std::string label = item.class_id >= 0 && item.class_id < static_cast<int>(backend->config.labels.size())
|
|
? backend->config.labels[static_cast<size_t>(item.class_id)] : "class_" + std::to_string(item.class_id);
|
|
std::snprintf(result.label, sizeof(result.label), "%s", label.c_str());
|
|
}
|
|
results->performance.postprocess_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - postprocess_start).count();
|
|
results->performance.total_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - total_start).count();
|
|
return YOLO_STATUS_OK;
|
|
}
|
|
}
|
|
|
|
Ort::Value* detection_value = nullptr;
|
|
Ort::Value* prototype_value = nullptr;
|
|
for (auto& output : output_values) {
|
|
if (!output.IsTensor()) continue;
|
|
const auto shape = output.GetTensorTypeAndShapeInfo().GetShape();
|
|
if (shape.size() == 3 && !detection_value) detection_value = &output;
|
|
if (shape.size() == 4 && !prototype_value) prototype_value = &output;
|
|
}
|
|
if (!detection_value) throw std::runtime_error("cannot find Ultralytics detection output");
|
|
const auto output_info = detection_value->GetTensorTypeAndShapeInfo();
|
|
const auto output_shape = output_info.GetShape();
|
|
if (output_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || output_shape.size() != 3)
|
|
throw std::runtime_error("Windows CPU Runtime requires FP32 [1,C,N] or [1,N,C] output");
|
|
const int64_t first_dimension = output_shape[1];
|
|
const int64_t second_dimension = output_shape[2];
|
|
// Usual Ultralytics output is [1, C, N] with N much larger than C.
|
|
// A valid one-candidate export has N == 1, however, so do not mistake
|
|
// [1, C, 1] for [1, N, C] merely because C is greater than N.
|
|
const bool channels_first = first_dimension > 4 &&
|
|
(second_dimension <= 4 || first_dimension <= second_dimension);
|
|
const int channels = static_cast<int>(channels_first ? first_dimension : second_dimension);
|
|
const int positions = static_cast<int>(channels_first ? second_dimension : first_dimension);
|
|
if (channels <= 4 || positions <= 0) throw std::runtime_error("invalid detection output shape");
|
|
|
|
int mask_channels = 0;
|
|
int prototype_height = 0;
|
|
int prototype_width = 0;
|
|
bool prototype_channels_first = true;
|
|
const float* prototypes = nullptr;
|
|
if (backend->config.task == YOLO_TASK_SEG) {
|
|
if (!prototype_value) throw std::runtime_error("cannot find Ultralytics segmentation prototype output");
|
|
const auto prototype_info = prototype_value->GetTensorTypeAndShapeInfo();
|
|
const auto prototype_shape = prototype_info.GetShape();
|
|
if (prototype_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || prototype_shape.size() != 4)
|
|
throw std::runtime_error("Windows CPU Runtime requires FP32 segmentation prototypes");
|
|
prototype_channels_first = prototype_shape[1] == 32 ||
|
|
(prototype_shape[1] > 0 && prototype_shape[1] <= 256 && prototype_shape[1] < prototype_shape[3]);
|
|
mask_channels = static_cast<int>(prototype_channels_first ? prototype_shape[1] : prototype_shape[3]);
|
|
prototype_height = static_cast<int>(prototype_channels_first ? prototype_shape[2] : prototype_shape[1]);
|
|
prototype_width = static_cast<int>(prototype_channels_first ? prototype_shape[3] : prototype_shape[2]);
|
|
if (mask_channels <= 0 || prototype_height <= 0 || prototype_width <= 0)
|
|
throw std::runtime_error("invalid segmentation prototype output shape");
|
|
prototypes = prototype_value->GetTensorData<float>();
|
|
}
|
|
|
|
const int geometry_channels = backend->config.task == YOLO_TASK_OBB ? 5 : 4;
|
|
int keypoint_count = 0;
|
|
int class_count = channels - geometry_channels - mask_channels;
|
|
if (backend->config.task == YOLO_TASK_POSE) {
|
|
class_count = backend->config.labels.empty() ? 1 : static_cast<int>(backend->config.labels.size());
|
|
const int keypoint_values = channels - 4 - class_count;
|
|
if (keypoint_values <= 0 || keypoint_values % 3 != 0)
|
|
throw std::runtime_error("pose output must contain x/y/score triplets after class scores");
|
|
keypoint_count = keypoint_values / 3;
|
|
if (keypoint_count > YOLO_BACKEND_MAX_KEYPOINTS)
|
|
throw std::runtime_error("pose output exceeds ABI keypoint capacity");
|
|
}
|
|
if (class_count <= 0) throw std::runtime_error("model output does not contain any class scores");
|
|
const float* output = detection_value->GetTensorData<float>();
|
|
const auto at = [&](int channel, int position) -> float {
|
|
return channels_first ? output[static_cast<size_t>(channel) * positions + position]
|
|
: output[static_cast<size_t>(position) * channels + channel];
|
|
};
|
|
std::vector<candidate> candidates;
|
|
for (int position = 0; position < positions; ++position) {
|
|
int class_id = 0;
|
|
float score = at(4, position);
|
|
for (int index = 1; index < class_count; ++index) {
|
|
const float value = at(4 + index, position);
|
|
if (value > score) { score = value; class_id = index; }
|
|
}
|
|
if (score < backend->confidence_threshold) continue;
|
|
const float width = std::clamp(at(2, position) / scale, 0.0f, static_cast<float>(image->width));
|
|
const float height = std::clamp(at(3, position) / scale, 0.0f, static_cast<float>(image->height));
|
|
candidate item{
|
|
class_id, score,
|
|
std::clamp((at(0, position) - pad_x) / scale - width * 0.5f, 0.0f, static_cast<float>(image->width)),
|
|
std::clamp((at(1, position) - pad_y) / scale - height * 0.5f, 0.0f, static_cast<float>(image->height)),
|
|
width, height };
|
|
if (backend->config.task == YOLO_TASK_OBB) item.angle = at(channels - 1, position);
|
|
if (backend->config.task == YOLO_TASK_SEG) {
|
|
item.mask_coefficients.resize(static_cast<size_t>(mask_channels));
|
|
for (int index = 0; index < mask_channels; ++index)
|
|
item.mask_coefficients[static_cast<size_t>(index)] = at(4 + class_count + index, position);
|
|
}
|
|
if (backend->config.task == YOLO_TASK_POSE) {
|
|
item.keypoints.resize(static_cast<size_t>(keypoint_count));
|
|
for (int index = 0; index < keypoint_count; ++index) {
|
|
const int offset = 4 + class_count + index * 3;
|
|
item.keypoints[static_cast<size_t>(index)] = {
|
|
std::clamp((at(offset, position) - pad_x) / scale, 0.0f, static_cast<float>(image->width)),
|
|
std::clamp((at(offset + 1, position) - pad_y) / scale, 0.0f, static_cast<float>(image->height)),
|
|
at(offset + 2, position)
|
|
};
|
|
}
|
|
}
|
|
candidates.push_back(std::move(item));
|
|
}
|
|
std::sort(candidates.begin(), candidates.end(), [](const candidate& a, const candidate& b) { return a.score > b.score; });
|
|
std::vector<candidate> selected_candidates;
|
|
for (const auto& item : candidates) {
|
|
if (selected_candidates.size() >= YOLO_BACKEND_MAX_RESULTS) break;
|
|
bool suppressed = false;
|
|
for (const auto& existing : selected_candidates) {
|
|
if ((backend->class_agnostic_nms || existing.class_id == item.class_id) &&
|
|
yolo_nms::overlap_iou(
|
|
backend->config.task, nms_box(existing), nms_box(item)) > backend->nms_threshold) {
|
|
suppressed = true;
|
|
break;
|
|
}
|
|
}
|
|
if (suppressed) continue;
|
|
selected_candidates.push_back(item);
|
|
}
|
|
|
|
backend->segmentation_masks.clear();
|
|
backend->segmentation_masks.reserve(selected_candidates.size());
|
|
const auto prototype_at = [&](int channel, int y, int x) -> float {
|
|
if (prototype_channels_first)
|
|
return prototypes[(static_cast<size_t>(channel) * prototype_height + y) * prototype_width + x];
|
|
return prototypes[(static_cast<size_t>(y) * prototype_width + x) * mask_channels + channel];
|
|
};
|
|
const auto mask_logit = [&](const candidate& item, float source_x, float source_y) -> float {
|
|
const float model_x = source_x * scale + pad_x;
|
|
const float model_y = source_y * scale + pad_y;
|
|
const float px = std::clamp((model_x + 0.5f) * prototype_width / backend->config.width - 0.5f,
|
|
0.0f, static_cast<float>(prototype_width - 1));
|
|
const float py = std::clamp((model_y + 0.5f) * prototype_height / backend->config.height - 0.5f,
|
|
0.0f, static_cast<float>(prototype_height - 1));
|
|
const int x0 = static_cast<int>(std::floor(px));
|
|
const int y0 = static_cast<int>(std::floor(py));
|
|
const int x1 = std::min(x0 + 1, prototype_width - 1);
|
|
const int y1 = std::min(y0 + 1, prototype_height - 1);
|
|
const float dx = px - x0;
|
|
const float dy = py - y0;
|
|
float value = 0.0f;
|
|
for (int channel = 0; channel < mask_channels; ++channel) {
|
|
const float top = prototype_at(channel, y0, x0) * (1.0f - dx) + prototype_at(channel, y0, x1) * dx;
|
|
const float bottom = prototype_at(channel, y1, x0) * (1.0f - dx) + prototype_at(channel, y1, x1) * dx;
|
|
value += item.mask_coefficients[static_cast<size_t>(channel)] * (top * (1.0f - dy) + bottom * dy);
|
|
}
|
|
return value;
|
|
};
|
|
|
|
for (const auto& item : selected_candidates) {
|
|
if (results->count >= YOLO_BACKEND_MAX_RESULTS) break;
|
|
int mask_min_x = image->width;
|
|
int mask_min_y = image->height;
|
|
int mask_max_x = -1;
|
|
int mask_max_y = -1;
|
|
std::vector<unsigned char> full_mask;
|
|
if (backend->config.task == YOLO_TASK_SEG) {
|
|
full_mask.assign(static_cast<size_t>(image->width) * image->height, 0);
|
|
const int start_x = std::clamp(static_cast<int>(std::floor(item.x)), 0, image->width);
|
|
const int start_y = std::clamp(static_cast<int>(std::floor(item.y)), 0, image->height);
|
|
const int end_x = std::clamp(static_cast<int>(std::ceil(item.x + item.width)), 0, image->width);
|
|
const int end_y = std::clamp(static_cast<int>(std::ceil(item.y + item.height)), 0, image->height);
|
|
for (int y = start_y; y < end_y; ++y) {
|
|
for (int x = start_x; x < end_x; ++x) {
|
|
if (mask_logit(item, x + 0.5f, y + 0.5f) <= 0.0f) continue;
|
|
full_mask[static_cast<size_t>(y) * image->width + x] = 255;
|
|
mask_min_x = std::min(mask_min_x, x);
|
|
mask_min_y = std::min(mask_min_y, y);
|
|
mask_max_x = std::max(mask_max_x, x);
|
|
mask_max_y = std::max(mask_max_y, y);
|
|
}
|
|
}
|
|
if (mask_max_x < mask_min_x || mask_max_y < mask_min_y) continue;
|
|
}
|
|
auto& result = results->results[results->count++];
|
|
result.task = backend->config.task;
|
|
result.class_id = item.class_id;
|
|
result.score = item.score;
|
|
result.bbox.x = item.x;
|
|
result.bbox.y = item.y;
|
|
result.bbox.width = item.width;
|
|
result.bbox.height = item.height;
|
|
if (backend->config.task == YOLO_TASK_OBB) {
|
|
result.bbox.angle = item.angle;
|
|
result.bbox.has_angle = 1;
|
|
}
|
|
if (backend->config.task == YOLO_TASK_POSE) {
|
|
result.keypoint_count = static_cast<int>(item.keypoints.size());
|
|
std::copy(item.keypoints.begin(), item.keypoints.end(), result.keypoints);
|
|
}
|
|
if (backend->config.task == YOLO_TASK_SEG) {
|
|
const int mask_width = mask_max_x - mask_min_x + 1;
|
|
const int mask_height = mask_max_y - mask_min_y + 1;
|
|
backend->segmentation_masks.emplace_back(static_cast<size_t>(mask_width) * mask_height);
|
|
auto& mask = backend->segmentation_masks.back();
|
|
for (int y = 0; y < mask_height; ++y)
|
|
std::copy_n(full_mask.data() + static_cast<size_t>(mask_min_y + y) * image->width + mask_min_x,
|
|
mask_width, mask.data() + static_cast<size_t>(y) * mask_width);
|
|
result.mask_bbox.x = static_cast<float>(mask_min_x);
|
|
result.mask_bbox.y = static_cast<float>(mask_min_y);
|
|
result.mask_bbox.width = static_cast<float>(mask_width);
|
|
result.mask_bbox.height = static_cast<float>(mask_height);
|
|
result.mask_data = mask.data();
|
|
result.mask_width = mask_width;
|
|
result.mask_height = mask_height;
|
|
result.mask_stride_bytes = mask_width;
|
|
result.mask_origin_x = mask_min_x;
|
|
result.mask_origin_y = mask_min_y;
|
|
}
|
|
const std::string label = item.class_id >= 0 && item.class_id < static_cast<int>(backend->config.labels.size())
|
|
? backend->config.labels[static_cast<size_t>(item.class_id)] : "class_" + std::to_string(item.class_id);
|
|
std::snprintf(result.label, sizeof(result.label), "%s", label.c_str());
|
|
}
|
|
results->performance.postprocess_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - postprocess_start).count();
|
|
results->performance.total_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - total_start).count();
|
|
return YOLO_STATUS_OK;
|
|
} catch (const std::exception& exception) {
|
|
backend->last_error = exception.what();
|
|
return YOLO_STATUS_BACKEND_FAILED;
|
|
}
|
|
}
|
|
|
|
extern "C" const char* yolo_backend_last_error(const yolo_runtime_backend_t* backend) {
|
|
return backend ? backend->last_error.c_str() : "backend is null";
|
|
}
|
|
|
|
extern "C" void yolo_backend_destroy(yolo_runtime_backend_t* backend) { delete backend; }
|
|
|