refactor open_vocab into yolo_onnx and yoloworld_onnx engines
Split the monolithic open_vocab engine into two focused engines: yolo_onnx (closed-vocabulary, any single-input YOLO ONNX) and yoloworld_onnx (open-vocabulary YOLO-World with CLIP encoder). Shared NMS/letterbox/fuse helpers extracted to yolo_helpers.h. Add export_yolo.py script and separate test binaries for each engine.
This commit is contained in:
cristhian aguilera
@@ -1,592 +0,0 @@
|
||||
#include "engines/open_vocab.h"
|
||||
#include "util/clip_tokenizer.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <onnxruntime_cxx_api.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <numeric>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
|
||||
// --- NMS helpers ---
|
||||
|
||||
struct RawDetection {
|
||||
float x1, y1, x2, y2;
|
||||
float score;
|
||||
int class_id;
|
||||
};
|
||||
|
||||
float iou(const RawDetection& a, const RawDetection& b) {
|
||||
float inter_x1 = std::max(a.x1, b.x1);
|
||||
float inter_y1 = std::max(a.y1, b.y1);
|
||||
float inter_x2 = std::min(a.x2, b.x2);
|
||||
float inter_y2 = std::min(a.y2, b.y2);
|
||||
float inter_w = std::max(0.0f, inter_x2 - inter_x1);
|
||||
float inter_h = std::max(0.0f, inter_y2 - inter_y1);
|
||||
float inter_area = inter_w * inter_h;
|
||||
float area_a = (a.x2 - a.x1) * (a.y2 - a.y1);
|
||||
float area_b = (b.x2 - b.x1) * (b.y2 - b.y1);
|
||||
float union_area = area_a + area_b - inter_area;
|
||||
return union_area > 0.0f ? inter_area / union_area : 0.0f;
|
||||
}
|
||||
|
||||
std::vector<RawDetection> nms(std::vector<RawDetection> dets, float threshold) {
|
||||
std::sort(dets.begin(), dets.end(),
|
||||
[](const RawDetection& a, const RawDetection& b) {
|
||||
return a.score > b.score;
|
||||
});
|
||||
std::vector<bool> suppressed(dets.size(), false);
|
||||
std::vector<RawDetection> result;
|
||||
for (size_t i = 0; i < dets.size(); i++) {
|
||||
if (suppressed[i]) continue;
|
||||
result.push_back(dets[i]);
|
||||
for (size_t j = i + 1; j < dets.size(); j++) {
|
||||
if (!suppressed[j] && iou(dets[i], dets[j]) > threshold) {
|
||||
suppressed[j] = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// --- Letterbox resize ---
|
||||
cv::Mat letterbox(const cv::Mat& src, int size, float& scale, float& pad_x, float& pad_y) {
|
||||
float scale_w = static_cast<float>(size) / src.cols;
|
||||
float scale_h = static_cast<float>(size) / src.rows;
|
||||
scale = std::min(scale_w, scale_h);
|
||||
int new_w = static_cast<int>(std::round(src.cols * scale));
|
||||
int new_h = static_cast<int>(std::round(src.rows * scale));
|
||||
pad_x = (size - new_w) / 2.0f;
|
||||
pad_y = (size - new_h) / 2.0f;
|
||||
|
||||
cv::Mat resized;
|
||||
cv::resize(src, resized, {new_w, new_h});
|
||||
|
||||
cv::Mat canvas(size, size, CV_8UC3, cv::Scalar(114, 114, 114));
|
||||
resized.copyTo(canvas(cv::Rect(
|
||||
static_cast<int>(pad_x), static_cast<int>(pad_y), new_w, new_h)));
|
||||
return canvas;
|
||||
}
|
||||
|
||||
// --- Point cloud fusion ---
|
||||
struct FusionResult {
|
||||
float cx, cy, cz;
|
||||
float depth_near, depth_far;
|
||||
bool valid = false;
|
||||
};
|
||||
|
||||
FusionResult fuse_bbox(const sl::Mat& pc, int x1, int y1, int x2, int y2) {
|
||||
FusionResult result{};
|
||||
int w = static_cast<int>(pc.getWidth());
|
||||
int h = static_cast<int>(pc.getHeight());
|
||||
x1 = std::max(x1, 0);
|
||||
y1 = std::max(y1, 0);
|
||||
x2 = std::min(x2, w - 1);
|
||||
y2 = std::min(y2, h - 1);
|
||||
|
||||
std::vector<float> xs, ys, zs;
|
||||
xs.reserve((x2 - x1) * (y2 - y1));
|
||||
ys.reserve(xs.capacity());
|
||||
zs.reserve(xs.capacity());
|
||||
|
||||
for (int y = y1; y <= y2; y++) {
|
||||
for (int x = x1; x <= x2; x++) {
|
||||
sl::float4 pt;
|
||||
pc.getValue(x, y, &pt);
|
||||
if (!std::isnan(pt.x) && !std::isnan(pt.y) && !std::isnan(pt.z) &&
|
||||
!std::isinf(pt.x) && !std::isinf(pt.y) && !std::isinf(pt.z)) {
|
||||
xs.push_back(pt.x);
|
||||
ys.push_back(pt.y);
|
||||
zs.push_back(pt.z);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (xs.empty()) return result;
|
||||
|
||||
std::sort(xs.begin(), xs.end());
|
||||
std::sort(ys.begin(), ys.end());
|
||||
std::sort(zs.begin(), zs.end());
|
||||
|
||||
size_t n = xs.size();
|
||||
size_t mid = n / 2;
|
||||
size_t p25 = n / 4;
|
||||
size_t p75 = 3 * n / 4;
|
||||
|
||||
result.cx = xs[mid];
|
||||
result.cy = ys[mid];
|
||||
result.cz = zs[mid];
|
||||
result.depth_near = zs[p25];
|
||||
result.depth_far = zs[p75];
|
||||
result.valid = true;
|
||||
return result;
|
||||
}
|
||||
|
||||
void backproject(float u, float v, float z, float fx, float fy, float cx, float cy,
|
||||
float& X, float& Y, float& Z) {
|
||||
Z = z;
|
||||
X = (u - cx) * z / fx;
|
||||
Y = (v - cy) * z / fy;
|
||||
}
|
||||
|
||||
// Split a comma-separated string into a vector of trimmed tokens
|
||||
std::vector<std::string> split_labels(const std::string& s) {
|
||||
std::vector<std::string> out;
|
||||
std::stringstream ss(s);
|
||||
std::string tok;
|
||||
while (std::getline(ss, tok, ',')) {
|
||||
size_t l = tok.find_first_not_of(" \t");
|
||||
size_t r = tok.find_last_not_of(" \t");
|
||||
if (l != std::string::npos)
|
||||
out.push_back(tok.substr(l, r - l + 1));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace engines {
|
||||
|
||||
// --- Impl ---
|
||||
|
||||
struct OpenVocab::Impl {
|
||||
// Visual model
|
||||
Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "open_vocab"};
|
||||
Ort::SessionOptions session_options;
|
||||
std::unique_ptr<Ort::Session> session;
|
||||
Ort::AllocatorWithDefaultOptions allocator;
|
||||
|
||||
std::string input_name; // visual model: image input
|
||||
std::string output_name; // visual model: detection output
|
||||
std::string text_input_name; // visual model: text_feats input (empty if single-input model)
|
||||
int num_classes = 0;
|
||||
bool has_text_input = false;
|
||||
|
||||
// Text encoder (optional; used when has_text_input == true)
|
||||
Ort::SessionOptions text_session_options;
|
||||
std::unique_ptr<Ort::Session> text_session;
|
||||
std::string text_enc_input_name; // "input_ids"
|
||||
std::string text_enc_output_name; // "text_embeds"
|
||||
|
||||
// Tokenizer + embedding cache
|
||||
util::ClipTokenizer tokenizer;
|
||||
std::vector<float> cached_text_feats; // (num_classes * 512) flat, row-major
|
||||
std::string cached_labels;
|
||||
|
||||
// Preprocessing + ONNX inference + NMS.
|
||||
// Takes a BGR image (cv::imread format); returns 2D detections in original image coords.
|
||||
std::vector<RawDetection> run_inference(const OpenVocabConfig& config,
|
||||
const cv::Mat& image_bgr);
|
||||
};
|
||||
|
||||
// --- Impl::run_inference ---
|
||||
|
||||
std::vector<RawDetection> OpenVocab::Impl::run_inference(
|
||||
const OpenVocabConfig& config, const cv::Mat& image_bgr)
|
||||
{
|
||||
if (!session) return {};
|
||||
if (has_text_input && cached_text_feats.empty()) return {};
|
||||
|
||||
int orig_w = image_bgr.cols;
|
||||
int orig_h = image_bgr.rows;
|
||||
|
||||
// BGR → RGB (YOLO expects RGB)
|
||||
cv::Mat rgb;
|
||||
cv::cvtColor(image_bgr, rgb, cv::COLOR_BGR2RGB);
|
||||
|
||||
float scale, pad_x, pad_y;
|
||||
cv::Mat lb = letterbox(rgb, config.input_size, scale, pad_x, pad_y);
|
||||
|
||||
cv::Mat lb_float;
|
||||
lb.convertTo(lb_float, CV_32F, 1.0f / 255.0f);
|
||||
|
||||
int sz = config.input_size;
|
||||
std::vector<float> input_data(3 * sz * sz);
|
||||
for (int c = 0; c < 3; c++) {
|
||||
for (int h = 0; h < sz; h++) {
|
||||
for (int w = 0; w < sz; w++) {
|
||||
input_data[c * sz * sz + h * sz + w] =
|
||||
lb_float.at<cv::Vec3f>(h, w)[c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
std::vector<int64_t> image_shape{1, 3, sz, sz};
|
||||
Ort::Value image_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem_info, input_data.data(), input_data.size(),
|
||||
image_shape.data(), image_shape.size());
|
||||
|
||||
const char* output_name_cstr = output_name.c_str();
|
||||
std::vector<Ort::Value> outputs;
|
||||
|
||||
try {
|
||||
if (has_text_input) {
|
||||
std::vector<int64_t> text_shape{1, num_classes, 512};
|
||||
Ort::Value text_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem_info,
|
||||
cached_text_feats.data(),
|
||||
cached_text_feats.size(),
|
||||
text_shape.data(), text_shape.size());
|
||||
|
||||
const char* input_names[2] = {
|
||||
input_name.c_str(),
|
||||
text_input_name.c_str()
|
||||
};
|
||||
Ort::Value input_tensors[2] = {
|
||||
std::move(image_tensor),
|
||||
std::move(text_tensor)
|
||||
};
|
||||
outputs = session->Run(
|
||||
Ort::RunOptions{nullptr},
|
||||
input_names, input_tensors, 2,
|
||||
&output_name_cstr, 1);
|
||||
} else {
|
||||
const char* input_name_cstr = input_name.c_str();
|
||||
outputs = session->Run(
|
||||
Ort::RunOptions{nullptr},
|
||||
&input_name_cstr, &image_tensor, 1,
|
||||
&output_name_cstr, 1);
|
||||
}
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "OpenVocab: inference failed: " << e.what() << std::endl;
|
||||
return {};
|
||||
}
|
||||
|
||||
// Post-process: output shape (1, 4+num_classes, num_anchors)
|
||||
auto out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
|
||||
if (out_shape.size() < 3) return {};
|
||||
|
||||
int num_anchors = static_cast<int>(out_shape[2]);
|
||||
int nc = num_classes;
|
||||
if (nc <= 0) return {};
|
||||
|
||||
const float* out = outputs[0].GetTensorData<float>();
|
||||
|
||||
std::vector<RawDetection> candidates;
|
||||
for (int a = 0; a < num_anchors; a++) {
|
||||
float cx_lb = out[0 * num_anchors + a];
|
||||
float cy_lb = out[1 * num_anchors + a];
|
||||
float w_lb = out[2 * num_anchors + a];
|
||||
float h_lb = out[3 * num_anchors + a];
|
||||
|
||||
float best_score = 0.0f;
|
||||
int best_class = 0;
|
||||
for (int c = 0; c < nc; c++) {
|
||||
float s = out[(4 + c) * num_anchors + a];
|
||||
if (s > best_score) { best_score = s; best_class = c; }
|
||||
}
|
||||
if (best_score < config.confidence_threshold) continue;
|
||||
|
||||
float x1_lb = cx_lb - w_lb / 2.0f;
|
||||
float y1_lb = cy_lb - h_lb / 2.0f;
|
||||
float x2_lb = cx_lb + w_lb / 2.0f;
|
||||
float y2_lb = cy_lb + h_lb / 2.0f;
|
||||
|
||||
float x1 = std::max(0.0f, std::min((x1_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
|
||||
candidates.push_back({x1, y1, x2, y2, best_score, best_class});
|
||||
}
|
||||
|
||||
return nms(std::move(candidates), config.nms_threshold);
|
||||
}
|
||||
|
||||
// --- Config ---
|
||||
|
||||
OpenVocabConfig OpenVocabConfig::from_env() {
|
||||
OpenVocabConfig c;
|
||||
c.model_path = util::get_env_str("ZED_ONNX_MODEL_PATH", c.model_path);
|
||||
c.confidence_threshold = util::get_env_float("ZED_ONNX_CONFIDENCE", c.confidence_threshold);
|
||||
c.nms_threshold = util::get_env_float("ZED_ONNX_NMS_THRESHOLD", c.nms_threshold);
|
||||
c.input_size = util::get_env_int("ZED_ONNX_INPUT_SIZE", c.input_size);
|
||||
c.labels = util::get_env_str("ZED_ONNX_LABELS", c.labels);
|
||||
c.text_encoder_path = util::get_env_str("ZED_ONNX_TEXT_ENCODER_PATH", c.text_encoder_path);
|
||||
c.clip_vocab_path = util::get_env_str("ZED_ONNX_CLIP_VOCAB", c.clip_vocab_path);
|
||||
c.clip_merges_path = util::get_env_str("ZED_ONNX_CLIP_MERGES", c.clip_merges_path);
|
||||
return c;
|
||||
}
|
||||
|
||||
// --- OpenVocab ---
|
||||
|
||||
OpenVocab::OpenVocab(const OpenVocabConfig& config)
|
||||
: config_(config), impl_(std::make_unique<Impl>()) {}
|
||||
|
||||
OpenVocab::~OpenVocab() = default;
|
||||
|
||||
void OpenVocab::enable(sl::Camera& /* camera */) {
|
||||
if (config_.model_path.empty()) {
|
||||
std::cerr << "OpenVocab: ZED_ONNX_MODEL_PATH is not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
// --- Load visual model ---
|
||||
OrtCUDAProviderOptions cuda_opts{};
|
||||
cuda_opts.device_id = 0;
|
||||
try {
|
||||
impl_->session_options.AppendExecutionProvider_CUDA(cuda_opts);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "OpenVocab: CUDA unavailable (" << e.what() << "), using CPU." << std::endl;
|
||||
}
|
||||
impl_->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
|
||||
std::cout << "OpenVocab: loading visual model (" << config_.model_path << ")..." << std::endl;
|
||||
try {
|
||||
impl_->session = std::make_unique<Ort::Session>(
|
||||
impl_->env, config_.model_path.c_str(), impl_->session_options);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "OpenVocab: failed to load visual model: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->input_name = impl_->session->GetInputNameAllocated(0, impl_->allocator).get();
|
||||
impl_->output_name = impl_->session->GetOutputNameAllocated(0, impl_->allocator).get();
|
||||
|
||||
size_t num_inputs = impl_->session->GetInputCount();
|
||||
impl_->has_text_input = (num_inputs >= 2);
|
||||
if (impl_->has_text_input) {
|
||||
impl_->text_input_name = impl_->session->GetInputNameAllocated(1, impl_->allocator).get();
|
||||
}
|
||||
|
||||
auto out_info = impl_->session->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo();
|
||||
auto shape = out_info.GetShape();
|
||||
if (shape.size() == 3 && shape[1] > 4) {
|
||||
impl_->num_classes = static_cast<int>(shape[1]) - 4;
|
||||
}
|
||||
|
||||
std::cout << "OpenVocab: visual model loaded."
|
||||
<< " Inputs=" << num_inputs
|
||||
<< " Classes=" << impl_->num_classes
|
||||
<< (impl_->has_text_input ? " (open-vocab)" : " (fixed-vocab)")
|
||||
<< std::endl;
|
||||
|
||||
// --- Load text encoder + tokenizer (only needed for open-vocab model) ---
|
||||
if (impl_->has_text_input) {
|
||||
if (config_.text_encoder_path.empty()) {
|
||||
std::cerr << "OpenVocab: ZED_ONNX_TEXT_ENCODER_PATH is not set "
|
||||
"(required for open-vocab model)." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->text_session_options.SetGraphOptimizationLevel(
|
||||
GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
|
||||
std::cout << "OpenVocab: loading text encoder (" << config_.text_encoder_path
|
||||
<< ")..." << std::endl;
|
||||
try {
|
||||
impl_->text_session = std::make_unique<Ort::Session>(
|
||||
impl_->env, config_.text_encoder_path.c_str(), impl_->text_session_options);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "OpenVocab: failed to load text encoder: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->text_enc_input_name =
|
||||
impl_->text_session->GetInputNameAllocated(0, impl_->allocator).get();
|
||||
impl_->text_enc_output_name =
|
||||
impl_->text_session->GetOutputNameAllocated(0, impl_->allocator).get();
|
||||
|
||||
if (!config_.clip_vocab_path.empty() && !config_.clip_merges_path.empty()) {
|
||||
try {
|
||||
impl_->tokenizer.load(config_.clip_vocab_path, config_.clip_merges_path);
|
||||
std::cout << "OpenVocab: tokenizer loaded." << std::endl;
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "OpenVocab: failed to load tokenizer: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
std::cerr << "OpenVocab: ZED_ONNX_CLIP_VOCAB / ZED_ONNX_CLIP_MERGES not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
if (!config_.labels.empty()) {
|
||||
encode_labels(config_.labels);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Encode a comma-separated label string into cached_text_feats.
|
||||
// Always fills exactly num_classes * 512 floats (pads with zeros if fewer labels).
|
||||
void OpenVocab::encode_labels(const std::string& labels_str) {
|
||||
if (!impl_->text_session || !impl_->tokenizer.is_loaded()) return;
|
||||
|
||||
auto label_vec = split_labels(labels_str);
|
||||
int N = impl_->num_classes;
|
||||
impl_->cached_text_feats.assign(static_cast<size_t>(N) * 512, 0.0f);
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
int encode_n = std::min(static_cast<int>(label_vec.size()), N);
|
||||
for (int i = 0; i < encode_n; i++) {
|
||||
auto tokens = impl_->tokenizer.tokenize(label_vec[i]);
|
||||
|
||||
std::vector<int64_t> input_shape{1, util::ClipTokenizer::SEQ_LEN};
|
||||
Ort::Value token_tensor = Ort::Value::CreateTensor<int64_t>(
|
||||
mem_info, tokens.data(), tokens.size(),
|
||||
input_shape.data(), input_shape.size());
|
||||
|
||||
const char* in = impl_->text_enc_input_name.c_str();
|
||||
const char* out = impl_->text_enc_output_name.c_str();
|
||||
std::vector<Ort::Value> outputs;
|
||||
try {
|
||||
outputs = impl_->text_session->Run(
|
||||
Ort::RunOptions{nullptr}, &in, &token_tensor, 1, &out, 1);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "OpenVocab: text encoder failed for \"" << label_vec[i]
|
||||
<< "\": " << e.what() << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
const float* emb = outputs[0].GetTensorData<float>();
|
||||
|
||||
// L2 normalize
|
||||
float norm = 0.0f;
|
||||
for (int j = 0; j < 512; j++) norm += emb[j] * emb[j];
|
||||
norm = std::sqrt(norm);
|
||||
float inv_norm = (norm > 1e-8f) ? 1.0f / norm : 0.0f;
|
||||
for (int j = 0; j < 512; j++)
|
||||
impl_->cached_text_feats[i * 512 + j] = emb[j] * inv_norm;
|
||||
}
|
||||
|
||||
impl_->cached_labels = labels_str;
|
||||
std::cout << "OpenVocab: encoded " << encode_n << " label(s): " << labels_str << std::endl;
|
||||
}
|
||||
|
||||
// --- detect(sl::Camera&) — production path with 3D fusion ---
|
||||
|
||||
void OpenVocab::detect(sl::Camera& camera) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
// Re-encode if labels changed (on-the-fly vocabulary switching)
|
||||
if (impl_->has_text_input) {
|
||||
const char* env_labels = std::getenv("ZED_ONNX_LABELS");
|
||||
std::string current_labels = env_labels ? env_labels : config_.labels;
|
||||
if (current_labels != impl_->cached_labels) {
|
||||
encode_labels(current_labels);
|
||||
}
|
||||
if (impl_->cached_text_feats.empty()) return;
|
||||
}
|
||||
|
||||
// Retrieve image and point cloud from camera
|
||||
sl::Mat img_mat;
|
||||
camera.retrieveImage(img_mat, sl::VIEW::LEFT);
|
||||
|
||||
sl::Mat pc_mat;
|
||||
camera.retrieveMeasure(pc_mat, sl::MEASURE::XYZRGBA);
|
||||
|
||||
sl::CameraInformation info = camera.getCameraInformation();
|
||||
auto calib = info.camera_configuration.calibration_parameters.left_cam;
|
||||
float fx = calib.fx, fy = calib.fy, cx = calib.cx, cy = calib.cy;
|
||||
|
||||
int orig_w = static_cast<int>(img_mat.getWidth());
|
||||
int orig_h = static_cast<int>(img_mat.getHeight());
|
||||
|
||||
// RGBA → BGR for run_inference
|
||||
cv::Mat rgba(orig_h, orig_w, CV_8UC4, img_mat.getPtr<sl::uchar1>());
|
||||
cv::Mat bgr;
|
||||
cv::cvtColor(rgba, bgr, cv::COLOR_RGBA2BGR);
|
||||
|
||||
auto kept = impl_->run_inference(config_, bgr);
|
||||
|
||||
// 3D fusion
|
||||
for (const auto& det : kept) {
|
||||
FusionResult fusion = fuse_bbox(pc_mat,
|
||||
static_cast<int>(det.x1), static_cast<int>(det.y1),
|
||||
static_cast<int>(det.x2), static_cast<int>(det.y2));
|
||||
|
||||
DetectedObject obj{};
|
||||
obj.label_id = det.class_id;
|
||||
obj.confidence = det.score * 100.0f;
|
||||
|
||||
if (fusion.valid) {
|
||||
obj.position[0] = fusion.cx;
|
||||
obj.position[1] = fusion.cy;
|
||||
obj.position[2] = fusion.cz;
|
||||
}
|
||||
|
||||
obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1;
|
||||
obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1;
|
||||
obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2;
|
||||
obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2;
|
||||
|
||||
if (fusion.valid) {
|
||||
float corners_u[4] = {det.x1, det.x2, det.x2, det.x1};
|
||||
float corners_v[4] = {det.y1, det.y1, det.y2, det.y2};
|
||||
float depths[2] = {fusion.depth_near, fusion.depth_far};
|
||||
for (int face = 0; face < 2; face++) {
|
||||
for (int c = 0; c < 4; c++) {
|
||||
float X, Y, Z;
|
||||
backproject(corners_u[c], corners_v[c], depths[face],
|
||||
fx, fy, cx, cy, X, Y, Z);
|
||||
int idx = (face * 4 + c) * 3;
|
||||
obj.bounding_box_3d[idx + 0] = X;
|
||||
obj.bounding_box_3d[idx + 1] = Y;
|
||||
obj.bounding_box_3d[idx + 2] = Z;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
detected_.push_back(obj);
|
||||
}
|
||||
}
|
||||
|
||||
// --- detect(const cv::Mat&) — image-only path, no 3D fusion ---
|
||||
|
||||
void OpenVocab::detect(const cv::Mat& image_bgr) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
// Re-encode if labels changed
|
||||
if (impl_->has_text_input) {
|
||||
const char* env_labels = std::getenv("ZED_ONNX_LABELS");
|
||||
std::string current_labels = env_labels ? env_labels : config_.labels;
|
||||
if (current_labels != impl_->cached_labels) {
|
||||
encode_labels(current_labels);
|
||||
}
|
||||
if (impl_->cached_text_feats.empty()) return;
|
||||
}
|
||||
|
||||
auto kept = impl_->run_inference(config_, image_bgr);
|
||||
|
||||
for (const auto& det : kept) {
|
||||
DetectedObject obj{};
|
||||
obj.label_id = det.class_id;
|
||||
obj.confidence = det.score * 100.0f;
|
||||
// position stays zero (no point cloud available)
|
||||
obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1;
|
||||
obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1;
|
||||
obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2;
|
||||
obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2;
|
||||
// bounding_box_3d stays zero
|
||||
detected_.push_back(obj);
|
||||
}
|
||||
}
|
||||
|
||||
void OpenVocab::set_labels(const std::string& labels) {
|
||||
config_.labels = labels;
|
||||
encode_labels(labels);
|
||||
}
|
||||
|
||||
const std::vector<DetectedObject>& OpenVocab::objects() const {
|
||||
return detected_;
|
||||
}
|
||||
|
||||
std::string OpenVocab::label_names() const {
|
||||
return impl_->cached_labels.empty() ? config_.labels : impl_->cached_labels;
|
||||
}
|
||||
|
||||
} // namespace engines
|
||||
@@ -0,0 +1,153 @@
|
||||
#pragma once
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace engines::detail {
|
||||
|
||||
// --- NMS ---
|
||||
|
||||
struct RawDetection {
|
||||
float x1, y1, x2, y2;
|
||||
float score;
|
||||
int class_id;
|
||||
};
|
||||
|
||||
inline float iou(const RawDetection& a, const RawDetection& b) {
|
||||
float inter_x1 = std::max(a.x1, b.x1);
|
||||
float inter_y1 = std::max(a.y1, b.y1);
|
||||
float inter_x2 = std::min(a.x2, b.x2);
|
||||
float inter_y2 = std::min(a.y2, b.y2);
|
||||
float inter_w = std::max(0.0f, inter_x2 - inter_x1);
|
||||
float inter_h = std::max(0.0f, inter_y2 - inter_y1);
|
||||
float inter_area = inter_w * inter_h;
|
||||
float area_a = (a.x2 - a.x1) * (a.y2 - a.y1);
|
||||
float area_b = (b.x2 - b.x1) * (b.y2 - b.y1);
|
||||
float union_area = area_a + area_b - inter_area;
|
||||
return union_area > 0.0f ? inter_area / union_area : 0.0f;
|
||||
}
|
||||
|
||||
inline std::vector<RawDetection> nms(std::vector<RawDetection> dets, float threshold) {
|
||||
std::sort(dets.begin(), dets.end(),
|
||||
[](const RawDetection& a, const RawDetection& b) {
|
||||
return a.score > b.score;
|
||||
});
|
||||
std::vector<bool> suppressed(dets.size(), false);
|
||||
std::vector<RawDetection> result;
|
||||
for (size_t i = 0; i < dets.size(); i++) {
|
||||
if (suppressed[i]) continue;
|
||||
result.push_back(dets[i]);
|
||||
for (size_t j = i + 1; j < dets.size(); j++) {
|
||||
if (!suppressed[j] && iou(dets[i], dets[j]) > threshold)
|
||||
suppressed[j] = true;
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// --- Letterbox resize ---
|
||||
|
||||
inline cv::Mat letterbox(const cv::Mat& src, int size,
|
||||
float& scale, float& pad_x, float& pad_y) {
|
||||
float scale_w = static_cast<float>(size) / src.cols;
|
||||
float scale_h = static_cast<float>(size) / src.rows;
|
||||
scale = std::min(scale_w, scale_h);
|
||||
int new_w = static_cast<int>(std::round(src.cols * scale));
|
||||
int new_h = static_cast<int>(std::round(src.rows * scale));
|
||||
pad_x = (size - new_w) / 2.0f;
|
||||
pad_y = (size - new_h) / 2.0f;
|
||||
|
||||
cv::Mat resized;
|
||||
cv::resize(src, resized, {new_w, new_h});
|
||||
|
||||
cv::Mat canvas(size, size, CV_8UC3, cv::Scalar(114, 114, 114));
|
||||
resized.copyTo(canvas(cv::Rect(
|
||||
static_cast<int>(pad_x), static_cast<int>(pad_y), new_w, new_h)));
|
||||
return canvas;
|
||||
}
|
||||
|
||||
// --- Point cloud fusion ---
|
||||
|
||||
struct FusionResult {
|
||||
float cx, cy, cz;
|
||||
float depth_near, depth_far;
|
||||
bool valid = false;
|
||||
};
|
||||
|
||||
inline FusionResult fuse_bbox(const sl::Mat& pc, int x1, int y1, int x2, int y2) {
|
||||
FusionResult result{};
|
||||
int w = static_cast<int>(pc.getWidth());
|
||||
int h = static_cast<int>(pc.getHeight());
|
||||
x1 = std::max(x1, 0);
|
||||
y1 = std::max(y1, 0);
|
||||
x2 = std::min(x2, w - 1);
|
||||
y2 = std::min(y2, h - 1);
|
||||
|
||||
std::vector<float> xs, ys, zs;
|
||||
xs.reserve((x2 - x1) * (y2 - y1));
|
||||
ys.reserve(xs.capacity());
|
||||
zs.reserve(xs.capacity());
|
||||
|
||||
for (int y = y1; y <= y2; y++) {
|
||||
for (int x = x1; x <= x2; x++) {
|
||||
sl::float4 pt;
|
||||
pc.getValue(x, y, &pt);
|
||||
if (!std::isnan(pt.x) && !std::isnan(pt.y) && !std::isnan(pt.z) &&
|
||||
!std::isinf(pt.x) && !std::isinf(pt.y) && !std::isinf(pt.z)) {
|
||||
xs.push_back(pt.x);
|
||||
ys.push_back(pt.y);
|
||||
zs.push_back(pt.z);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (xs.empty()) return result;
|
||||
|
||||
std::sort(xs.begin(), xs.end());
|
||||
std::sort(ys.begin(), ys.end());
|
||||
std::sort(zs.begin(), zs.end());
|
||||
|
||||
size_t n = xs.size();
|
||||
size_t mid = n / 2;
|
||||
size_t p25 = n / 4;
|
||||
size_t p75 = 3 * n / 4;
|
||||
|
||||
result.cx = xs[mid];
|
||||
result.cy = ys[mid];
|
||||
result.cz = zs[mid];
|
||||
result.depth_near = zs[p25];
|
||||
result.depth_far = zs[p75];
|
||||
result.valid = true;
|
||||
return result;
|
||||
}
|
||||
|
||||
inline void backproject(float u, float v, float z,
|
||||
float fx, float fy, float cx, float cy,
|
||||
float& X, float& Y, float& Z) {
|
||||
Z = z;
|
||||
X = (u - cx) * z / fx;
|
||||
Y = (v - cy) * z / fy;
|
||||
}
|
||||
|
||||
// --- Label parsing ---
|
||||
|
||||
inline std::vector<std::string> split_labels(const std::string& s) {
|
||||
std::vector<std::string> out;
|
||||
std::stringstream ss(s);
|
||||
std::string tok;
|
||||
while (std::getline(ss, tok, ',')) {
|
||||
size_t l = tok.find_first_not_of(" \t");
|
||||
size_t r = tok.find_last_not_of(" \t");
|
||||
if (l != std::string::npos)
|
||||
out.push_back(tok.substr(l, r - l + 1));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
} // namespace engines::detail
|
||||
@@ -0,0 +1,373 @@
|
||||
#include "engines/yolo_onnx.h"
|
||||
#include "engines/yolo_helpers.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <onnxruntime_cxx_api.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <cmath>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace engines {
|
||||
|
||||
using namespace detail;
|
||||
|
||||
// --- Impl ---
|
||||
|
||||
struct YoloOnnx::Impl {
|
||||
Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "yolo_onnx"};
|
||||
Ort::SessionOptions session_options;
|
||||
std::unique_ptr<Ort::Session> session;
|
||||
Ort::AllocatorWithDefaultOptions allocator;
|
||||
|
||||
std::string input_name;
|
||||
std::string output_name;
|
||||
int num_classes = 0;
|
||||
YoloOutputLayout resolved_layout = YoloOutputLayout::NC_FIRST;
|
||||
|
||||
std::vector<RawDetection> run_inference(const YoloOnnxConfig& config,
|
||||
const cv::Mat& image_bgr);
|
||||
};
|
||||
|
||||
// --- Impl::run_inference ---
|
||||
|
||||
std::vector<RawDetection> YoloOnnx::Impl::run_inference(
|
||||
const YoloOnnxConfig& config, const cv::Mat& image_bgr)
|
||||
{
|
||||
if (!session) return {};
|
||||
|
||||
int orig_w = image_bgr.cols;
|
||||
int orig_h = image_bgr.rows;
|
||||
|
||||
cv::Mat rgb;
|
||||
cv::cvtColor(image_bgr, rgb, cv::COLOR_BGR2RGB);
|
||||
|
||||
float scale, pad_x, pad_y;
|
||||
cv::Mat lb = letterbox(rgb, config.input_size, scale, pad_x, pad_y);
|
||||
|
||||
cv::Mat lb_float;
|
||||
lb.convertTo(lb_float, CV_32F, 1.0f / 255.0f);
|
||||
|
||||
int sz = config.input_size;
|
||||
std::vector<float> input_data(3 * sz * sz);
|
||||
for (int c = 0; c < 3; c++) {
|
||||
for (int h = 0; h < sz; h++) {
|
||||
for (int w = 0; w < sz; w++) {
|
||||
input_data[c * sz * sz + h * sz + w] =
|
||||
lb_float.at<cv::Vec3f>(h, w)[c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
std::vector<int64_t> image_shape{1, 3, sz, sz};
|
||||
Ort::Value image_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem_info, input_data.data(), input_data.size(),
|
||||
image_shape.data(), image_shape.size());
|
||||
|
||||
const char* input_name_cstr = input_name.c_str();
|
||||
const char* output_name_cstr = output_name.c_str();
|
||||
|
||||
std::vector<Ort::Value> outputs;
|
||||
try {
|
||||
outputs = session->Run(
|
||||
Ort::RunOptions{nullptr},
|
||||
&input_name_cstr, &image_tensor, 1,
|
||||
&output_name_cstr, 1);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloOnnx: inference failed: " << e.what() << std::endl;
|
||||
return {};
|
||||
}
|
||||
|
||||
auto out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
|
||||
if (out_shape.size() < 3) return {};
|
||||
|
||||
const float* out = outputs[0].GetTensorData<float>();
|
||||
int nc = num_classes;
|
||||
if (nc <= 0) return {};
|
||||
|
||||
std::vector<RawDetection> candidates;
|
||||
|
||||
if (resolved_layout == YoloOutputLayout::NC_FIRST) {
|
||||
// (1, 4+N, anchors) — YOLOv8, YOLOv11
|
||||
int num_anchors = static_cast<int>(out_shape[2]);
|
||||
for (int a = 0; a < num_anchors; a++) {
|
||||
float cx_lb = out[0 * num_anchors + a];
|
||||
float cy_lb = out[1 * num_anchors + a];
|
||||
float w_lb = out[2 * num_anchors + a];
|
||||
float h_lb = out[3 * num_anchors + a];
|
||||
|
||||
float best_score = 0.0f;
|
||||
int best_class = 0;
|
||||
for (int c = 0; c < nc; c++) {
|
||||
float s = out[(4 + c) * num_anchors + a];
|
||||
if (s > best_score) { best_score = s; best_class = c; }
|
||||
}
|
||||
if (best_score < config.confidence_threshold) continue;
|
||||
|
||||
float x1_lb = cx_lb - w_lb / 2.0f;
|
||||
float y1_lb = cy_lb - h_lb / 2.0f;
|
||||
float x2_lb = cx_lb + w_lb / 2.0f;
|
||||
float y2_lb = cy_lb + h_lb / 2.0f;
|
||||
|
||||
float x1 = std::max(0.0f, std::min((x1_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
candidates.push_back({x1, y1, x2, y2, best_score, best_class});
|
||||
}
|
||||
} else {
|
||||
// (1, anchors, 4+N) or (1, anchors, 5+N)
|
||||
int num_anchors = static_cast<int>(out_shape[1]);
|
||||
int row_stride = static_cast<int>(out_shape[2]);
|
||||
int cls_offset = (resolved_layout == YoloOutputLayout::NC_LAST_OBJ) ? 5 : 4;
|
||||
|
||||
for (int a = 0; a < num_anchors; a++) {
|
||||
const float* row = out + a * row_stride;
|
||||
float cx_lb = row[0], cy_lb = row[1], w_lb = row[2], h_lb = row[3];
|
||||
float obj_conf = (resolved_layout == YoloOutputLayout::NC_LAST_OBJ) ? row[4] : 1.0f;
|
||||
|
||||
float best_score = 0.0f;
|
||||
int best_class = 0;
|
||||
for (int c = 0; c < nc; c++) {
|
||||
float s = obj_conf * row[cls_offset + c];
|
||||
if (s > best_score) { best_score = s; best_class = c; }
|
||||
}
|
||||
if (best_score < config.confidence_threshold) continue;
|
||||
|
||||
float x1_lb = cx_lb - w_lb / 2.0f;
|
||||
float y1_lb = cy_lb - h_lb / 2.0f;
|
||||
float x2_lb = cx_lb + w_lb / 2.0f;
|
||||
float y2_lb = cy_lb + h_lb / 2.0f;
|
||||
|
||||
float x1 = std::max(0.0f, std::min((x1_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
candidates.push_back({x1, y1, x2, y2, best_score, best_class});
|
||||
}
|
||||
}
|
||||
|
||||
return nms(std::move(candidates), config.nms_threshold);
|
||||
}
|
||||
|
||||
// --- Config ---
|
||||
|
||||
YoloOnnxConfig YoloOnnxConfig::from_env() {
|
||||
YoloOnnxConfig c;
|
||||
c.model_path = util::get_env_str("ZED_ONNX_MODEL_PATH", c.model_path);
|
||||
c.confidence_threshold = util::get_env_float("ZED_ONNX_CONFIDENCE", c.confidence_threshold);
|
||||
c.nms_threshold = util::get_env_float("ZED_ONNX_NMS_THRESHOLD", c.nms_threshold);
|
||||
c.input_size = util::get_env_int("ZED_ONNX_INPUT_SIZE", c.input_size);
|
||||
c.labels = util::get_env_str("ZED_ONNX_LABELS", c.labels);
|
||||
c.labels_file = util::get_env_str("ZED_ONNX_LABELS_FILE", c.labels_file);
|
||||
|
||||
std::string fmt = util::get_env_str("ZED_ONNX_OUTPUT_FORMAT", "");
|
||||
if (fmt == "nc_first" || fmt == "yolov8" || fmt == "yolov11")
|
||||
c.output_layout = YoloOutputLayout::NC_FIRST;
|
||||
else if (fmt == "nc_last")
|
||||
c.output_layout = YoloOutputLayout::NC_LAST;
|
||||
else if (fmt == "nc_last_obj" || fmt == "yolov5")
|
||||
c.output_layout = YoloOutputLayout::NC_LAST_OBJ;
|
||||
|
||||
return c;
|
||||
}
|
||||
|
||||
// --- YoloOnnx ---
|
||||
|
||||
YoloOnnx::YoloOnnx(const YoloOnnxConfig& config)
|
||||
: config_(config), impl_(std::make_unique<Impl>()) {}
|
||||
|
||||
YoloOnnx::~YoloOnnx() = default;
|
||||
|
||||
void YoloOnnx::enable(sl::Camera& /* camera */) {
|
||||
if (config_.model_path.empty()) {
|
||||
std::cerr << "YoloOnnx: ZED_ONNX_MODEL_PATH is not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
OrtCUDAProviderOptions cuda_opts{};
|
||||
cuda_opts.device_id = 0;
|
||||
try {
|
||||
impl_->session_options.AppendExecutionProvider_CUDA(cuda_opts);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloOnnx: CUDA unavailable (" << e.what() << "), using CPU." << std::endl;
|
||||
}
|
||||
impl_->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
|
||||
std::cout << "YoloOnnx: loading model (" << config_.model_path << ")..." << std::endl;
|
||||
try {
|
||||
impl_->session = std::make_unique<Ort::Session>(
|
||||
impl_->env, config_.model_path.c_str(), impl_->session_options);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloOnnx: failed to load model: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->input_name = impl_->session->GetInputNameAllocated(0, impl_->allocator).get();
|
||||
impl_->output_name = impl_->session->GetOutputNameAllocated(0, impl_->allocator).get();
|
||||
|
||||
size_t num_inputs = impl_->session->GetInputCount();
|
||||
if (num_inputs >= 2) {
|
||||
std::cerr << "YoloOnnx: model has " << num_inputs << " inputs — "
|
||||
"this looks like a YOLO-World model. Use ZED_ENGINE=yoloworld_onnx instead."
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
auto out_info = impl_->session->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo();
|
||||
auto shape = out_info.GetShape();
|
||||
if (shape.size() == 3) {
|
||||
int64_t dim1 = shape[1];
|
||||
int64_t dim2 = shape[2];
|
||||
|
||||
YoloOutputLayout resolved;
|
||||
if (config_.output_layout == YoloOutputLayout::AUTO) {
|
||||
// dim1 < dim2 → NC_FIRST (e.g. 84 vs 8400 for YOLOv8/11)
|
||||
// dim1 > dim2 → NC_LAST_OBJ (e.g. 25200 vs 85 for YOLOv5)
|
||||
resolved = (dim1 < dim2) ? YoloOutputLayout::NC_FIRST
|
||||
: YoloOutputLayout::NC_LAST_OBJ;
|
||||
} else {
|
||||
resolved = config_.output_layout;
|
||||
}
|
||||
impl_->resolved_layout = resolved;
|
||||
|
||||
if (resolved == YoloOutputLayout::NC_FIRST) {
|
||||
impl_->num_classes = static_cast<int>(dim1) - 4;
|
||||
} else if (resolved == YoloOutputLayout::NC_LAST) {
|
||||
impl_->num_classes = static_cast<int>(dim2) - 4;
|
||||
} else { // NC_LAST_OBJ
|
||||
impl_->num_classes = static_cast<int>(dim2) - 5;
|
||||
}
|
||||
}
|
||||
|
||||
if (!config_.labels_file.empty() && config_.labels.empty()) {
|
||||
std::ifstream lf(config_.labels_file);
|
||||
if (!lf) {
|
||||
std::cerr << "YoloOnnx: cannot open ZED_ONNX_LABELS_FILE: "
|
||||
<< config_.labels_file << std::endl;
|
||||
} else {
|
||||
std::string line, joined;
|
||||
while (std::getline(lf, line)) {
|
||||
size_t l = line.find_first_not_of(" \t\r");
|
||||
size_t r = line.find_last_not_of(" \t\r");
|
||||
if (l == std::string::npos) continue;
|
||||
if (!joined.empty()) joined += ',';
|
||||
joined += line.substr(l, r - l + 1);
|
||||
}
|
||||
if (!joined.empty()) {
|
||||
config_.labels = joined;
|
||||
std::cout << "YoloOnnx: loaded labels from file ("
|
||||
<< config_.labels_file << ")" << std::endl;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::cout << "YoloOnnx: model loaded."
|
||||
<< " Classes=" << impl_->num_classes
|
||||
<< " Layout=" << (impl_->resolved_layout == YoloOutputLayout::NC_FIRST ? "NC_FIRST"
|
||||
: impl_->resolved_layout == YoloOutputLayout::NC_LAST ? "NC_LAST"
|
||||
: "NC_LAST_OBJ")
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
void YoloOnnx::detect(sl::Camera& camera) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
sl::Mat img_mat;
|
||||
camera.retrieveImage(img_mat, sl::VIEW::LEFT);
|
||||
|
||||
sl::Mat pc_mat;
|
||||
camera.retrieveMeasure(pc_mat, sl::MEASURE::XYZRGBA);
|
||||
|
||||
sl::CameraInformation info = camera.getCameraInformation();
|
||||
auto calib = info.camera_configuration.calibration_parameters.left_cam;
|
||||
float fx = calib.fx, fy = calib.fy, cx = calib.cx, cy = calib.cy;
|
||||
|
||||
int orig_w = static_cast<int>(img_mat.getWidth());
|
||||
int orig_h = static_cast<int>(img_mat.getHeight());
|
||||
|
||||
cv::Mat rgba(orig_h, orig_w, CV_8UC4, img_mat.getPtr<sl::uchar1>());
|
||||
cv::Mat bgr;
|
||||
cv::cvtColor(rgba, bgr, cv::COLOR_RGBA2BGR);
|
||||
|
||||
auto kept = impl_->run_inference(config_, bgr);
|
||||
|
||||
for (const auto& det : kept) {
|
||||
FusionResult fusion = fuse_bbox(pc_mat,
|
||||
static_cast<int>(det.x1), static_cast<int>(det.y1),
|
||||
static_cast<int>(det.x2), static_cast<int>(det.y2));
|
||||
|
||||
DetectedObject obj{};
|
||||
obj.label_id = det.class_id;
|
||||
obj.confidence = det.score * 100.0f;
|
||||
|
||||
if (fusion.valid) {
|
||||
obj.position[0] = fusion.cx;
|
||||
obj.position[1] = fusion.cy;
|
||||
obj.position[2] = fusion.cz;
|
||||
}
|
||||
|
||||
obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1;
|
||||
obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1;
|
||||
obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2;
|
||||
obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2;
|
||||
|
||||
if (fusion.valid) {
|
||||
float corners_u[4] = {det.x1, det.x2, det.x2, det.x1};
|
||||
float corners_v[4] = {det.y1, det.y1, det.y2, det.y2};
|
||||
float depths[2] = {fusion.depth_near, fusion.depth_far};
|
||||
for (int face = 0; face < 2; face++) {
|
||||
for (int c = 0; c < 4; c++) {
|
||||
float X, Y, Z;
|
||||
backproject(corners_u[c], corners_v[c], depths[face],
|
||||
fx, fy, cx, cy, X, Y, Z);
|
||||
int idx = (face * 4 + c) * 3;
|
||||
obj.bounding_box_3d[idx + 0] = X;
|
||||
obj.bounding_box_3d[idx + 1] = Y;
|
||||
obj.bounding_box_3d[idx + 2] = Z;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
detected_.push_back(obj);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloOnnx::detect(const cv::Mat& image_bgr) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
auto kept = impl_->run_inference(config_, image_bgr);
|
||||
|
||||
for (const auto& det : kept) {
|
||||
DetectedObject obj{};
|
||||
obj.label_id = det.class_id;
|
||||
obj.confidence = det.score * 100.0f;
|
||||
obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1;
|
||||
obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1;
|
||||
obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2;
|
||||
obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2;
|
||||
detected_.push_back(obj);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloOnnx::set_labels(const std::string& labels) {
|
||||
config_.labels = labels;
|
||||
}
|
||||
|
||||
const std::vector<DetectedObject>& YoloOnnx::objects() const {
|
||||
return detected_;
|
||||
}
|
||||
|
||||
std::string YoloOnnx::label_names() const {
|
||||
return config_.labels;
|
||||
}
|
||||
|
||||
} // namespace engines
|
||||
@@ -0,0 +1,49 @@
|
||||
#pragma once
|
||||
|
||||
#include "engines/engine.h"
|
||||
|
||||
#include <opencv2/core.hpp>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
namespace engines {
|
||||
|
||||
enum class YoloOutputLayout {
|
||||
AUTO, // detect from output tensor shape at load time
|
||||
NC_FIRST, // (1, 4+N, anchors) — YOLOv8, YOLOv11
|
||||
NC_LAST, // (1, anchors, 4+N) — rows-first, no objectness
|
||||
NC_LAST_OBJ // (1, anchors, 5+N) — YOLOv5, rows-first with objectness
|
||||
};
|
||||
|
||||
struct YoloOnnxConfig {
|
||||
std::string model_path; // ZED_ONNX_MODEL_PATH (required)
|
||||
float confidence_threshold = 0.25f; // ZED_ONNX_CONFIDENCE
|
||||
float nms_threshold = 0.45f; // ZED_ONNX_NMS_THRESHOLD
|
||||
int input_size = 640; // ZED_ONNX_INPUT_SIZE
|
||||
std::string labels; // ZED_ONNX_LABELS (comma-separated)
|
||||
std::string labels_file; // ZED_ONNX_LABELS_FILE (one class name per line)
|
||||
YoloOutputLayout output_layout = YoloOutputLayout::AUTO;
|
||||
|
||||
static YoloOnnxConfig from_env();
|
||||
};
|
||||
|
||||
class YoloOnnx : public Engine {
|
||||
public:
|
||||
explicit YoloOnnx(const YoloOnnxConfig& config);
|
||||
~YoloOnnx() override;
|
||||
|
||||
void enable(sl::Camera& camera) override;
|
||||
void detect(sl::Camera& camera) override;
|
||||
void detect(const cv::Mat& image_bgr); // image-only; no 3D fusion
|
||||
void set_labels(const std::string& labels) override;
|
||||
const std::vector<DetectedObject>& objects() const override;
|
||||
std::string label_names() const override;
|
||||
|
||||
private:
|
||||
YoloOnnxConfig config_;
|
||||
struct Impl;
|
||||
std::unique_ptr<Impl> impl_;
|
||||
std::vector<DetectedObject> detected_;
|
||||
};
|
||||
|
||||
} // namespace engines
|
||||
@@ -0,0 +1,415 @@
|
||||
#include "engines/yoloworld_onnx.h"
|
||||
#include "engines/yolo_helpers.h"
|
||||
#include "util/clip_tokenizer.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <onnxruntime_cxx_api.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace engines {
|
||||
|
||||
using namespace detail;
|
||||
|
||||
// --- Impl ---
|
||||
|
||||
struct YoloWorldOnnx::Impl {
|
||||
// Visual model
|
||||
Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "yoloworld_onnx"};
|
||||
Ort::SessionOptions session_options;
|
||||
std::unique_ptr<Ort::Session> session;
|
||||
Ort::AllocatorWithDefaultOptions allocator;
|
||||
|
||||
std::string input_name; // image input
|
||||
std::string text_input_name; // text_feats input
|
||||
std::string output_name;
|
||||
int num_classes = 0;
|
||||
|
||||
// Text encoder
|
||||
Ort::SessionOptions text_session_options;
|
||||
std::unique_ptr<Ort::Session> text_session;
|
||||
std::string text_enc_input_name; // "input_ids"
|
||||
std::string text_enc_output_name; // "text_embeds"
|
||||
|
||||
// Tokenizer + embedding cache
|
||||
util::ClipTokenizer tokenizer;
|
||||
std::vector<float> cached_text_feats; // (num_classes * 512) flat, row-major
|
||||
std::string cached_labels;
|
||||
|
||||
std::vector<RawDetection> run_inference(const YoloWorldOnnxConfig& config,
|
||||
const cv::Mat& image_bgr);
|
||||
};
|
||||
|
||||
// --- Impl::run_inference ---
|
||||
|
||||
std::vector<RawDetection> YoloWorldOnnx::Impl::run_inference(
|
||||
const YoloWorldOnnxConfig& config, const cv::Mat& image_bgr)
|
||||
{
|
||||
if (!session) return {};
|
||||
if (cached_text_feats.empty()) return {};
|
||||
|
||||
int orig_w = image_bgr.cols;
|
||||
int orig_h = image_bgr.rows;
|
||||
|
||||
cv::Mat rgb;
|
||||
cv::cvtColor(image_bgr, rgb, cv::COLOR_BGR2RGB);
|
||||
|
||||
float scale, pad_x, pad_y;
|
||||
cv::Mat lb = letterbox(rgb, config.input_size, scale, pad_x, pad_y);
|
||||
|
||||
cv::Mat lb_float;
|
||||
lb.convertTo(lb_float, CV_32F, 1.0f / 255.0f);
|
||||
|
||||
int sz = config.input_size;
|
||||
std::vector<float> input_data(3 * sz * sz);
|
||||
for (int c = 0; c < 3; c++) {
|
||||
for (int h = 0; h < sz; h++) {
|
||||
for (int w = 0; w < sz; w++) {
|
||||
input_data[c * sz * sz + h * sz + w] =
|
||||
lb_float.at<cv::Vec3f>(h, w)[c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
std::vector<int64_t> image_shape{1, 3, sz, sz};
|
||||
Ort::Value image_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem_info, input_data.data(), input_data.size(),
|
||||
image_shape.data(), image_shape.size());
|
||||
|
||||
std::vector<int64_t> text_shape{1, num_classes, 512};
|
||||
Ort::Value text_tensor = Ort::Value::CreateTensor<float>(
|
||||
mem_info,
|
||||
cached_text_feats.data(),
|
||||
cached_text_feats.size(),
|
||||
text_shape.data(), text_shape.size());
|
||||
|
||||
const char* input_names[2] = {input_name.c_str(), text_input_name.c_str()};
|
||||
Ort::Value input_tensors[2] = {std::move(image_tensor), std::move(text_tensor)};
|
||||
const char* output_name_cstr = output_name.c_str();
|
||||
|
||||
std::vector<Ort::Value> outputs;
|
||||
try {
|
||||
outputs = session->Run(
|
||||
Ort::RunOptions{nullptr},
|
||||
input_names, input_tensors, 2,
|
||||
&output_name_cstr, 1);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: inference failed: " << e.what() << std::endl;
|
||||
return {};
|
||||
}
|
||||
|
||||
// Output: (1, 4+num_classes, num_anchors) — NC_FIRST only
|
||||
auto out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
|
||||
if (out_shape.size() < 3) return {};
|
||||
|
||||
int num_anchors = static_cast<int>(out_shape[2]);
|
||||
int nc = num_classes;
|
||||
if (nc <= 0) return {};
|
||||
|
||||
const float* out = outputs[0].GetTensorData<float>();
|
||||
|
||||
std::vector<RawDetection> candidates;
|
||||
for (int a = 0; a < num_anchors; a++) {
|
||||
float cx_lb = out[0 * num_anchors + a];
|
||||
float cy_lb = out[1 * num_anchors + a];
|
||||
float w_lb = out[2 * num_anchors + a];
|
||||
float h_lb = out[3 * num_anchors + a];
|
||||
|
||||
float best_score = 0.0f;
|
||||
int best_class = 0;
|
||||
for (int c = 0; c < nc; c++) {
|
||||
float s = out[(4 + c) * num_anchors + a];
|
||||
if (s > best_score) { best_score = s; best_class = c; }
|
||||
}
|
||||
if (best_score < config.confidence_threshold) continue;
|
||||
|
||||
float x1_lb = cx_lb - w_lb / 2.0f;
|
||||
float y1_lb = cy_lb - h_lb / 2.0f;
|
||||
float x2_lb = cx_lb + w_lb / 2.0f;
|
||||
float y2_lb = cy_lb + h_lb / 2.0f;
|
||||
|
||||
float x1 = std::max(0.0f, std::min((x1_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
|
||||
float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
|
||||
|
||||
candidates.push_back({x1, y1, x2, y2, best_score, best_class});
|
||||
}
|
||||
|
||||
return nms(std::move(candidates), config.nms_threshold);
|
||||
}
|
||||
|
||||
// --- Config ---
|
||||
|
||||
YoloWorldOnnxConfig YoloWorldOnnxConfig::from_env() {
|
||||
YoloWorldOnnxConfig c;
|
||||
c.model_path = util::get_env_str("ZED_ONNX_MODEL_PATH", c.model_path);
|
||||
c.confidence_threshold = util::get_env_float("ZED_ONNX_CONFIDENCE", c.confidence_threshold);
|
||||
c.nms_threshold = util::get_env_float("ZED_ONNX_NMS_THRESHOLD", c.nms_threshold);
|
||||
c.input_size = util::get_env_int("ZED_ONNX_INPUT_SIZE", c.input_size);
|
||||
c.labels = util::get_env_str("ZED_ONNX_LABELS", c.labels);
|
||||
c.text_encoder_path = util::get_env_str("ZED_ONNX_TEXT_ENCODER_PATH", c.text_encoder_path);
|
||||
c.clip_vocab_path = util::get_env_str("ZED_ONNX_CLIP_VOCAB", c.clip_vocab_path);
|
||||
c.clip_merges_path = util::get_env_str("ZED_ONNX_CLIP_MERGES", c.clip_merges_path);
|
||||
return c;
|
||||
}
|
||||
|
||||
// --- YoloWorldOnnx ---
|
||||
|
||||
YoloWorldOnnx::YoloWorldOnnx(const YoloWorldOnnxConfig& config)
|
||||
: config_(config), impl_(std::make_unique<Impl>()) {}
|
||||
|
||||
YoloWorldOnnx::~YoloWorldOnnx() = default;
|
||||
|
||||
void YoloWorldOnnx::enable(sl::Camera& /* camera */) {
|
||||
if (config_.model_path.empty()) {
|
||||
std::cerr << "YoloWorldOnnx: ZED_ONNX_MODEL_PATH is not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
// --- Load visual model ---
|
||||
OrtCUDAProviderOptions cuda_opts{};
|
||||
cuda_opts.device_id = 0;
|
||||
try {
|
||||
impl_->session_options.AppendExecutionProvider_CUDA(cuda_opts);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: CUDA unavailable (" << e.what() << "), using CPU." << std::endl;
|
||||
}
|
||||
impl_->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
|
||||
std::cout << "YoloWorldOnnx: loading visual model (" << config_.model_path << ")..." << std::endl;
|
||||
try {
|
||||
impl_->session = std::make_unique<Ort::Session>(
|
||||
impl_->env, config_.model_path.c_str(), impl_->session_options);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: failed to load visual model: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
size_t num_inputs = impl_->session->GetInputCount();
|
||||
if (num_inputs < 2) {
|
||||
std::cerr << "YoloWorldOnnx: model has only " << num_inputs << " input(s) — "
|
||||
"expected a two-input YOLO-World model. Use ZED_ENGINE=yolo_onnx for single-input models."
|
||||
<< std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->input_name = impl_->session->GetInputNameAllocated(0, impl_->allocator).get();
|
||||
impl_->text_input_name = impl_->session->GetInputNameAllocated(1, impl_->allocator).get();
|
||||
impl_->output_name = impl_->session->GetOutputNameAllocated(0, impl_->allocator).get();
|
||||
|
||||
auto out_info = impl_->session->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo();
|
||||
auto shape = out_info.GetShape();
|
||||
if (shape.size() == 3 && shape[1] > 4) {
|
||||
impl_->num_classes = static_cast<int>(shape[1]) - 4;
|
||||
}
|
||||
|
||||
std::cout << "YoloWorldOnnx: visual model loaded."
|
||||
<< " Classes=" << impl_->num_classes
|
||||
<< std::endl;
|
||||
|
||||
// --- Load text encoder + tokenizer ---
|
||||
if (config_.text_encoder_path.empty()) {
|
||||
std::cerr << "YoloWorldOnnx: ZED_ONNX_TEXT_ENCODER_PATH is not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->text_session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||||
|
||||
std::cout << "YoloWorldOnnx: loading text encoder (" << config_.text_encoder_path << ")..." << std::endl;
|
||||
try {
|
||||
impl_->text_session = std::make_unique<Ort::Session>(
|
||||
impl_->env, config_.text_encoder_path.c_str(), impl_->text_session_options);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: failed to load text encoder: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
impl_->text_enc_input_name =
|
||||
impl_->text_session->GetInputNameAllocated(0, impl_->allocator).get();
|
||||
impl_->text_enc_output_name =
|
||||
impl_->text_session->GetOutputNameAllocated(0, impl_->allocator).get();
|
||||
|
||||
if (!config_.clip_vocab_path.empty() && !config_.clip_merges_path.empty()) {
|
||||
try {
|
||||
impl_->tokenizer.load(config_.clip_vocab_path, config_.clip_merges_path);
|
||||
std::cout << "YoloWorldOnnx: tokenizer loaded." << std::endl;
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: failed to load tokenizer: " << e.what() << std::endl;
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
std::cerr << "YoloWorldOnnx: ZED_ONNX_CLIP_VOCAB / ZED_ONNX_CLIP_MERGES not set." << std::endl;
|
||||
return;
|
||||
}
|
||||
|
||||
if (!config_.labels.empty()) {
|
||||
encode_labels(config_.labels);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloWorldOnnx::encode_labels(const std::string& labels_str) {
|
||||
if (!impl_->text_session || !impl_->tokenizer.is_loaded()) return;
|
||||
|
||||
auto label_vec = split_labels(labels_str);
|
||||
int N = impl_->num_classes;
|
||||
impl_->cached_text_feats.assign(static_cast<size_t>(N) * 512, 0.0f);
|
||||
|
||||
auto mem_info = Ort::MemoryInfo::CreateCpu(
|
||||
OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);
|
||||
|
||||
int encode_n = std::min(static_cast<int>(label_vec.size()), N);
|
||||
for (int i = 0; i < encode_n; i++) {
|
||||
auto tokens = impl_->tokenizer.tokenize(label_vec[i]);
|
||||
|
||||
std::vector<int64_t> input_shape{1, util::ClipTokenizer::SEQ_LEN};
|
||||
Ort::Value token_tensor = Ort::Value::CreateTensor<int64_t>(
|
||||
mem_info, tokens.data(), tokens.size(),
|
||||
input_shape.data(), input_shape.size());
|
||||
|
||||
const char* in = impl_->text_enc_input_name.c_str();
|
||||
const char* out = impl_->text_enc_output_name.c_str();
|
||||
std::vector<Ort::Value> outputs;
|
||||
try {
|
||||
outputs = impl_->text_session->Run(
|
||||
Ort::RunOptions{nullptr}, &in, &token_tensor, 1, &out, 1);
|
||||
} catch (const Ort::Exception& e) {
|
||||
std::cerr << "YoloWorldOnnx: text encoder failed for \"" << label_vec[i]
|
||||
<< "\": " << e.what() << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
const float* emb = outputs[0].GetTensorData<float>();
|
||||
|
||||
float norm = 0.0f;
|
||||
for (int j = 0; j < 512; j++) norm += emb[j] * emb[j];
|
||||
norm = std::sqrt(norm);
|
||||
float inv_norm = (norm > 1e-8f) ? 1.0f / norm : 0.0f;
|
||||
for (int j = 0; j < 512; j++)
|
||||
impl_->cached_text_feats[i * 512 + j] = emb[j] * inv_norm;
|
||||
}
|
||||
|
||||
impl_->cached_labels = labels_str;
|
||||
std::cout << "YoloWorldOnnx: encoded " << encode_n << " label(s): " << labels_str << std::endl;
|
||||
}
|
||||
|
||||
void YoloWorldOnnx::detect(sl::Camera& camera) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
// Re-encode if labels changed (on-the-fly vocabulary switching)
|
||||
const char* env_labels = std::getenv("ZED_ONNX_LABELS");
|
||||
std::string current_labels = env_labels ? env_labels : config_.labels;
|
||||
if (current_labels != impl_->cached_labels) {
|
||||
encode_labels(current_labels);
|
||||
}
|
||||
if (impl_->cached_text_feats.empty()) return;
|
||||
|
||||
sl::Mat img_mat;
|
||||
camera.retrieveImage(img_mat, sl::VIEW::LEFT);
|
||||
|
||||
sl::Mat pc_mat;
|
||||
camera.retrieveMeasure(pc_mat, sl::MEASURE::XYZRGBA);
|
||||
|
||||
sl::CameraInformation info = camera.getCameraInformation();
|
||||
auto calib = info.camera_configuration.calibration_parameters.left_cam;
|
||||
float fx = calib.fx, fy = calib.fy, cx = calib.cx, cy = calib.cy;
|
||||
|
||||
int orig_w = static_cast<int>(img_mat.getWidth());
|
||||
int orig_h = static_cast<int>(img_mat.getHeight());
|
||||
|
||||
cv::Mat rgba(orig_h, orig_w, CV_8UC4, img_mat.getPtr<sl::uchar1>());
|
||||
cv::Mat bgr;
|
||||
cv::cvtColor(rgba, bgr, cv::COLOR_RGBA2BGR);
|
||||
|
||||
auto kept = impl_->run_inference(config_, bgr);
|
||||
|
||||
for (const auto& det : kept) {
|
||||
FusionResult fusion = fuse_bbox(pc_mat,
|
||||
static_cast<int>(det.x1), static_cast<int>(det.y1),
|
||||
static_cast<int>(det.x2), static_cast<int>(det.y2));
|
||||
|
||||
DetectedObject obj{};
|
||||
obj.label_id = det.class_id;
|
||||
obj.confidence = det.score * 100.0f;
|
||||
|
||||
if (fusion.valid) {
|
||||
obj.position[0] = fusion.cx;
|
||||
obj.position[1] = fusion.cy;
|
||||
obj.position[2] = fusion.cz;
|
||||
}
|
||||
|
||||
obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1;
|
||||
obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1;
|
||||
obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2;
|
||||
obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2;
|
||||
|
||||
if (fusion.valid) {
|
||||
float corners_u[4] = {det.x1, det.x2, det.x2, det.x1};
|
||||
float corners_v[4] = {det.y1, det.y1, det.y2, det.y2};
|
||||
float depths[2] = {fusion.depth_near, fusion.depth_far};
|
||||
for (int face = 0; face < 2; face++) {
|
||||
for (int c = 0; c < 4; c++) {
|
||||
float X, Y, Z;
|
||||
backproject(corners_u[c], corners_v[c], depths[face],
|
||||
fx, fy, cx, cy, X, Y, Z);
|
||||
int idx = (face * 4 + c) * 3;
|
||||
obj.bounding_box_3d[idx + 0] = X;
|
||||
obj.bounding_box_3d[idx + 1] = Y;
|
||||
obj.bounding_box_3d[idx + 2] = Z;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
detected_.push_back(obj);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloWorldOnnx::detect(const cv::Mat& image_bgr) {
|
||||
detected_.clear();
|
||||
if (!impl_->session) return;
|
||||
|
||||
const char* env_labels = std::getenv("ZED_ONNX_LABELS");
|
||||
std::string current_labels = env_labels ? env_labels : config_.labels;
|
||||
if (current_labels != impl_->cached_labels) {
|
||||
encode_labels(current_labels);
|
||||
}
|
||||
if (impl_->cached_text_feats.empty()) return;
|
||||
|
||||
auto kept = impl_->run_inference(config_, image_bgr);
|
||||
|
||||
for (const auto& det : kept) {
|
||||
DetectedObject obj{};
|
||||
obj.label_id = det.class_id;
|
||||
obj.confidence = det.score * 100.0f;
|
||||
obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1;
|
||||
obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1;
|
||||
obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2;
|
||||
obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2;
|
||||
detected_.push_back(obj);
|
||||
}
|
||||
}
|
||||
|
||||
void YoloWorldOnnx::set_labels(const std::string& labels) {
|
||||
config_.labels = labels;
|
||||
encode_labels(labels);
|
||||
}
|
||||
|
||||
const std::vector<DetectedObject>& YoloWorldOnnx::objects() const {
|
||||
return detected_;
|
||||
}
|
||||
|
||||
std::string YoloWorldOnnx::label_names() const {
|
||||
return impl_->cached_labels.empty() ? config_.labels : impl_->cached_labels;
|
||||
}
|
||||
|
||||
} // namespace engines
|
||||
@@ -8,35 +8,35 @@
|
||||
|
||||
namespace engines {
|
||||
|
||||
struct OpenVocabConfig {
|
||||
struct YoloWorldOnnxConfig {
|
||||
std::string model_path; // ZED_ONNX_MODEL_PATH (required)
|
||||
float confidence_threshold = 0.25f; // ZED_ONNX_CONFIDENCE
|
||||
float nms_threshold = 0.45f; // ZED_ONNX_NMS_THRESHOLD
|
||||
int input_size = 640; // ZED_ONNX_INPUT_SIZE
|
||||
std::string labels; // ZED_ONNX_LABELS (comma-separated)
|
||||
|
||||
// Open-vocabulary text encoder (required when model has two inputs)
|
||||
// CLIP text encoder (required)
|
||||
std::string text_encoder_path; // ZED_ONNX_TEXT_ENCODER_PATH
|
||||
std::string clip_vocab_path; // ZED_ONNX_CLIP_VOCAB
|
||||
std::string clip_merges_path; // ZED_ONNX_CLIP_MERGES
|
||||
|
||||
static OpenVocabConfig from_env();
|
||||
static YoloWorldOnnxConfig from_env();
|
||||
};
|
||||
|
||||
class OpenVocab : public Engine {
|
||||
class YoloWorldOnnx : public Engine {
|
||||
public:
|
||||
explicit OpenVocab(const OpenVocabConfig& config);
|
||||
~OpenVocab() override;
|
||||
explicit YoloWorldOnnx(const YoloWorldOnnxConfig& config);
|
||||
~YoloWorldOnnx() override;
|
||||
|
||||
void enable(sl::Camera& camera) override;
|
||||
void detect(sl::Camera& camera) override;
|
||||
void detect(const cv::Mat& image_bgr); // image-only; no 3D fusion
|
||||
void detect(const cv::Mat& image_bgr); // image-only; no 3D fusion
|
||||
void set_labels(const std::string& labels) override;
|
||||
const std::vector<DetectedObject>& objects() const override;
|
||||
std::string label_names() const override;
|
||||
|
||||
private:
|
||||
OpenVocabConfig config_;
|
||||
YoloWorldOnnxConfig config_;
|
||||
struct Impl;
|
||||
std::unique_ptr<Impl> impl_;
|
||||
std::vector<DetectedObject> detected_;
|
||||
+8
-4
@@ -1,5 +1,6 @@
|
||||
#include "engines/engine.h"
|
||||
#include "engines/open_vocab.h"
|
||||
#include "engines/yolo_onnx.h"
|
||||
#include "engines/yoloworld_onnx.h"
|
||||
#include "engines/zed_native.h"
|
||||
#include "util/env.h"
|
||||
#include "zed/camera.h"
|
||||
@@ -74,9 +75,12 @@ int main() {
|
||||
|
||||
std::string engine_name = util::get_env_str("ZED_ENGINE", "native");
|
||||
std::unique_ptr<engines::Engine> detector;
|
||||
if (engine_name == "open_vocab") {
|
||||
auto cfg = engines::OpenVocabConfig::from_env();
|
||||
detector = std::make_unique<engines::OpenVocab>(cfg);
|
||||
if (engine_name == "yolo_onnx" || engine_name == "onnx") {
|
||||
auto cfg = engines::YoloOnnxConfig::from_env();
|
||||
detector = std::make_unique<engines::YoloOnnx>(cfg);
|
||||
} else if (engine_name == "yoloworld_onnx" || engine_name == "open_vocab") {
|
||||
auto cfg = engines::YoloWorldOnnxConfig::from_env();
|
||||
detector = std::make_unique<engines::YoloWorldOnnx>(cfg);
|
||||
} else {
|
||||
auto cfg = engines::DetectionConfig::from_env();
|
||||
detector = std::make_unique<engines::ZedNative>(cfg);
|
||||
|
||||
@@ -1,24 +1,22 @@
|
||||
// Standalone test for the OpenVocab detection engine.
|
||||
// No dora framework, no live ZED camera — loads images from disk.
|
||||
// Standalone test for the YoloOnnx (closed-vocabulary) detection engine.
|
||||
// No dora framework, no live ZED camera — loads an image from disk.
|
||||
//
|
||||
// Usage:
|
||||
// bin/test_open_vocab [yolo_image] [yoloworld_image]
|
||||
// yolo_image default: test/assets/yolo.png
|
||||
// yoloworld_image default: test/assets/yoloworld.png
|
||||
// bin/test_yolo_onnx [image_path]
|
||||
// image_path default: test/assets/yolo.png
|
||||
//
|
||||
// Uses the same env vars as the main binary:
|
||||
// ZED_ONNX_MODEL_PATH path to YOLO ONNX model (required)
|
||||
// ZED_ONNX_LABELS comma-separated labels (required)
|
||||
// ZED_ONNX_CONFIDENCE detection threshold (default 0.25)
|
||||
// ZED_ONNX_NMS_THRESHOLD NMS IoU threshold (default 0.45)
|
||||
// ZED_ONNX_INPUT_SIZE model input size (default 640)
|
||||
// ZED_ONNX_TEXT_ENCODER_PATH clip text encoder ONNX (open-vocab only)
|
||||
// ZED_ONNX_CLIP_VOCAB clip_vocab.txt path (open-vocab only)
|
||||
// ZED_ONNX_CLIP_MERGES clip_merges.txt path (open-vocab only)
|
||||
// ZED_ONNX_WORLD_MODEL_PATH alternative ONNX model for the second image;
|
||||
// falls back to ZED_ONNX_MODEL_PATH if not set
|
||||
// Required env vars:
|
||||
// ZED_ONNX_MODEL_PATH path to YOLO ONNX model
|
||||
//
|
||||
// Optional env vars:
|
||||
// ZED_ONNX_LABELS comma-separated class names
|
||||
// ZED_ONNX_LABELS_FILE path to newline-separated class names file
|
||||
// ZED_ONNX_CONFIDENCE detection threshold (default 0.25)
|
||||
// ZED_ONNX_NMS_THRESHOLD NMS IoU threshold (default 0.45)
|
||||
// ZED_ONNX_INPUT_SIZE model input size (default 640)
|
||||
// ZED_ONNX_OUTPUT_FORMAT nc_first|nc_last|nc_last_obj|yolov5|yolov8|yolov11|auto
|
||||
|
||||
#include "engines/open_vocab.h"
|
||||
#include "engines/yolo_onnx.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
@@ -61,19 +59,28 @@ void draw_detections(cv::Mat& img,
|
||||
|
||||
char buf[64];
|
||||
std::snprintf(buf, sizeof(buf), " %.2f", obj.confidence / 100.0f);
|
||||
std::string text = name + buf;
|
||||
|
||||
cv::rectangle(img, {x1, y1}, {x2, y2}, {0, 200, 0}, 2);
|
||||
cv::putText(img, text, {x1, std::max(y1 - 6, 14)},
|
||||
cv::putText(img, name + buf, {x1, std::max(y1 - 6, 14)},
|
||||
cv::FONT_HERSHEY_SIMPLEX, 0.55, {0, 200, 0}, 2, cv::LINE_AA);
|
||||
}
|
||||
}
|
||||
|
||||
bool run_test(engines::OpenVocab& engine, const std::string& image_path) {
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
std::string image_path = (argc > 1) ? argv[1] : "test/assets/yolo.png";
|
||||
|
||||
auto config = engines::YoloOnnxConfig::from_env();
|
||||
|
||||
sl::Camera dummy_cam;
|
||||
engines::YoloOnnx engine(config);
|
||||
engine.enable(dummy_cam);
|
||||
|
||||
cv::Mat img = cv::imread(image_path);
|
||||
if (img.empty()) {
|
||||
std::cerr << "Error: cannot read image: " << image_path << std::endl;
|
||||
return false;
|
||||
return 1;
|
||||
}
|
||||
|
||||
engine.detect(img);
|
||||
@@ -101,42 +108,6 @@ bool run_test(engines::OpenVocab& engine, const std::string& image_path) {
|
||||
std::string result_path = (parent / (stem.string() + "_result.png")).string();
|
||||
cv::imwrite(result_path, annotated);
|
||||
std::cout << "Saved: " << result_path << std::endl;
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
std::string yolo_path = (argc > 1) ? argv[1] : "test/assets/yolo.png";
|
||||
std::string yoloworld_path = (argc > 2) ? argv[2] : "test/assets/yoloworld.png";
|
||||
|
||||
auto config = engines::OpenVocabConfig::from_env();
|
||||
|
||||
// enable() ignores this parameter for OpenVocab
|
||||
sl::Camera dummy_cam;
|
||||
|
||||
int failures = 0;
|
||||
|
||||
// --- Test 1: yolo image ---
|
||||
{
|
||||
std::cout << "\n=== Test 1: " << yolo_path << " ===" << std::endl;
|
||||
engines::OpenVocab engine(config);
|
||||
engine.enable(dummy_cam);
|
||||
if (!run_test(engine, yolo_path)) ++failures;
|
||||
}
|
||||
|
||||
// --- Test 2: yoloworld image (optionally with a different model) ---
|
||||
{
|
||||
std::cout << "\n=== Test 2: " << yoloworld_path << " ===" << std::endl;
|
||||
engines::OpenVocabConfig config2 = config;
|
||||
const char* world_model = std::getenv("ZED_ONNX_WORLD_MODEL_PATH");
|
||||
if (world_model && *world_model)
|
||||
config2.model_path = world_model;
|
||||
|
||||
engines::OpenVocab engine(config2);
|
||||
engine.enable(dummy_cam);
|
||||
if (!run_test(engine, yoloworld_path)) ++failures;
|
||||
}
|
||||
|
||||
return failures > 0 ? 1 : 0;
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
// Standalone test for the YoloWorldOnnx (open-vocabulary) detection engine.
|
||||
// No dora framework, no live ZED camera — loads an image from disk.
|
||||
//
|
||||
// Usage:
|
||||
// bin/test_yoloworld_onnx [image_path]
|
||||
// image_path default: test/assets/yoloworld.png
|
||||
//
|
||||
// Required env vars:
|
||||
// ZED_ONNX_MODEL_PATH path to YOLO-World ONNX model (two-input)
|
||||
// ZED_ONNX_TEXT_ENCODER_PATH path to clip_text_encoder.onnx
|
||||
// ZED_ONNX_CLIP_VOCAB path to clip_vocab.txt
|
||||
// ZED_ONNX_CLIP_MERGES path to clip_merges.txt
|
||||
// ZED_ONNX_LABELS comma-separated class labels
|
||||
//
|
||||
// Optional env vars:
|
||||
// ZED_ONNX_CONFIDENCE detection threshold (default 0.25)
|
||||
// ZED_ONNX_NMS_THRESHOLD NMS IoU threshold (default 0.45)
|
||||
// ZED_ONNX_INPUT_SIZE model input size (default 640)
|
||||
|
||||
#include "engines/yoloworld_onnx.h"
|
||||
#include "util/env.h"
|
||||
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <sl/Camera.hpp>
|
||||
|
||||
#include <filesystem>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
namespace {
|
||||
|
||||
std::vector<std::string> parse_csv(const std::string& s) {
|
||||
std::vector<std::string> out;
|
||||
std::istringstream ss(s);
|
||||
std::string tok;
|
||||
while (std::getline(ss, tok, ',')) {
|
||||
size_t l = tok.find_first_not_of(" \t");
|
||||
size_t r = tok.find_last_not_of(" \t");
|
||||
if (l != std::string::npos)
|
||||
out.push_back(tok.substr(l, r - l + 1));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
void draw_detections(cv::Mat& img,
|
||||
const std::vector<engines::DetectedObject>& objects,
|
||||
const std::vector<std::string>& labels) {
|
||||
for (const auto& obj : objects) {
|
||||
int x1 = static_cast<int>(obj.bounding_box_2d[0]);
|
||||
int y1 = static_cast<int>(obj.bounding_box_2d[1]);
|
||||
int x2 = static_cast<int>(obj.bounding_box_2d[4]);
|
||||
int y2 = static_cast<int>(obj.bounding_box_2d[5]);
|
||||
|
||||
std::string name = (obj.label_id >= 0 && obj.label_id < static_cast<int>(labels.size()))
|
||||
? labels[obj.label_id] : std::to_string(obj.label_id);
|
||||
|
||||
char buf[64];
|
||||
std::snprintf(buf, sizeof(buf), " %.2f", obj.confidence / 100.0f);
|
||||
|
||||
cv::rectangle(img, {x1, y1}, {x2, y2}, {0, 200, 0}, 2);
|
||||
cv::putText(img, name + buf, {x1, std::max(y1 - 6, 14)},
|
||||
cv::FONT_HERSHEY_SIMPLEX, 0.55, {0, 200, 0}, 2, cv::LINE_AA);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
std::string image_path = (argc > 1) ? argv[1] : "test/assets/yoloworld.png";
|
||||
|
||||
auto config = engines::YoloWorldOnnxConfig::from_env();
|
||||
|
||||
sl::Camera dummy_cam;
|
||||
engines::YoloWorldOnnx engine(config);
|
||||
engine.enable(dummy_cam);
|
||||
|
||||
cv::Mat img = cv::imread(image_path);
|
||||
if (img.empty()) {
|
||||
std::cerr << "Error: cannot read image: " << image_path << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
engine.detect(img);
|
||||
const auto& objects = engine.objects();
|
||||
auto labels = parse_csv(engine.label_names());
|
||||
|
||||
std::cout << "[" << image_path << "] " << objects.size() << " detection(s):" << std::endl;
|
||||
for (const auto& obj : objects) {
|
||||
std::string name = (obj.label_id >= 0 && obj.label_id < static_cast<int>(labels.size()))
|
||||
? labels[obj.label_id] : std::to_string(obj.label_id);
|
||||
std::cout << " " << name
|
||||
<< " conf=" << obj.confidence / 100.0f
|
||||
<< " bbox=["
|
||||
<< obj.bounding_box_2d[0] << "," << obj.bounding_box_2d[1]
|
||||
<< " -> "
|
||||
<< obj.bounding_box_2d[4] << "," << obj.bounding_box_2d[5]
|
||||
<< "]" << std::endl;
|
||||
}
|
||||
|
||||
cv::Mat annotated = img.clone();
|
||||
draw_detections(annotated, objects, labels);
|
||||
|
||||
fs::path stem = fs::path(image_path).stem();
|
||||
fs::path parent = fs::path(image_path).parent_path();
|
||||
std::string result_path = (parent / (stem.string() + "_result.png")).string();
|
||||
cv::imwrite(result_path, annotated);
|
||||
std::cout << "Saved: " << result_path << std::endl;
|
||||
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user