cristhian aguilera
44eec29b83
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.
299 lines
12 KiB
C++
299 lines
12 KiB
C++
#include "engines/engine.h"
|
|
#include "engines/yolo_onnx.h"
|
|
#include "engines/yoloworld_onnx.h"
|
|
#include "engines/zed_native.h"
|
|
#include "util/env.h"
|
|
#include "zed/camera.h"
|
|
|
|
#include <dora-node-api.h>
|
|
#include <opencv2/opencv.hpp>
|
|
|
|
#include <chrono>
|
|
#include <cstring>
|
|
#include <initializer_list>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace {
|
|
|
|
constexpr int FIELDS_PER_OBJECT = 37;
|
|
|
|
enum class PcMode { FULL, HALF, ROI };
|
|
|
|
rust::Vec<int64_t> make_shape(int height, int width, int channels) {
|
|
rust::Vec<int64_t> shape;
|
|
shape.push_back(height);
|
|
shape.push_back(width);
|
|
shape.push_back(channels);
|
|
return shape;
|
|
}
|
|
|
|
rust::Vec<double> make_distortion(const std::vector<double>& distortion) {
|
|
rust::Vec<double> values;
|
|
for (double v : distortion) {
|
|
values.push_back(v);
|
|
}
|
|
return values;
|
|
}
|
|
|
|
rust::Vec<int64_t> make_vector_int(std::initializer_list<int64_t> values) {
|
|
rust::Vec<int64_t> out;
|
|
for (int64_t v : values) {
|
|
out.push_back(v);
|
|
}
|
|
return out;
|
|
}
|
|
|
|
std::vector<float> serialize_objects(const std::vector<engines::DetectedObject>& objects) {
|
|
std::vector<float> buffer(objects.size() * FIELDS_PER_OBJECT);
|
|
for (size_t i = 0; i < objects.size(); i++) {
|
|
float* row = buffer.data() + i * FIELDS_PER_OBJECT;
|
|
const auto& obj = objects[i];
|
|
row[0] = static_cast<float>(obj.label_id);
|
|
row[1] = obj.confidence;
|
|
row[2] = obj.position[0];
|
|
row[3] = obj.position[1];
|
|
row[4] = obj.position[2];
|
|
std::memcpy(row + 5, obj.bounding_box_2d, sizeof(obj.bounding_box_2d));
|
|
std::memcpy(row + 13, obj.bounding_box_3d, sizeof(obj.bounding_box_3d));
|
|
}
|
|
return buffer;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main() {
|
|
auto dora_node = init_dora_node();
|
|
|
|
auto config = zed::Config::from_env();
|
|
zed::Camera camera(config);
|
|
camera.open();
|
|
camera.warmup();
|
|
auto calib = camera.calibration();
|
|
|
|
std::string engine_name = util::get_env_str("ZED_ENGINE", "native");
|
|
std::unique_ptr<engines::Engine> detector;
|
|
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);
|
|
}
|
|
detector->enable(camera.handle());
|
|
|
|
std::string pc_mode_str = util::get_env_str("ZED_PC_MODE", "FULL");
|
|
PcMode pc_mode = (pc_mode_str == "HALF") ? PcMode::HALF
|
|
: (pc_mode_str == "ROI") ? PcMode::ROI
|
|
: PcMode::FULL;
|
|
|
|
int64_t frame_id = 0;
|
|
while (true) {
|
|
auto event = dora_node.events->next();
|
|
auto type = event_type(event);
|
|
|
|
if (type == DoraEventType::Stop || type == DoraEventType::AllInputsClosed) {
|
|
break;
|
|
}
|
|
|
|
if (type != DoraEventType::Input) {
|
|
continue;
|
|
}
|
|
|
|
auto input = event_as_input(std::move(event));
|
|
std::string input_id = std::string(input.id);
|
|
|
|
if (input_id == "update_vocabulary") {
|
|
if (!input.data.empty()) {
|
|
std::string labels(
|
|
reinterpret_cast<const char*>(input.data.data()),
|
|
input.data.size());
|
|
detector->set_labels(labels);
|
|
std::cout << "Vocabulary updated: " << labels << std::endl;
|
|
}
|
|
continue;
|
|
}
|
|
|
|
if (!camera.grab()) {
|
|
std::cerr << "Grab failed; attempting to reconnect..." << std::endl;
|
|
camera.close();
|
|
camera.open();
|
|
calib = camera.calibration();
|
|
detector->enable(camera.handle());
|
|
continue;
|
|
}
|
|
|
|
detector->detect(camera.handle());
|
|
|
|
sl::Mat& image_mat = camera.image();
|
|
sl::Mat& point_cloud_mat = camera.point_cloud();
|
|
|
|
int width = image_mat.getWidth();
|
|
int height = image_mat.getHeight();
|
|
|
|
int64_t current_frame_id = frame_id++;
|
|
int64_t timestamp_ns = static_cast<int64_t>(
|
|
std::chrono::duration_cast<std::chrono::nanoseconds>(
|
|
std::chrono::steady_clock::now().time_since_epoch()
|
|
).count()
|
|
);
|
|
|
|
// --- image_bgr ---
|
|
cv::Mat image_rgba(height, width, CV_8UC4, image_mat.getPtr<sl::uchar1>());
|
|
cv::Mat image_bgr;
|
|
cv::cvtColor(image_rgba, image_bgr, cv::COLOR_BGRA2BGR);
|
|
|
|
const auto image_size = static_cast<size_t>(image_bgr.total() * image_bgr.elemSize());
|
|
rust::Slice<const uint8_t> image_slice{
|
|
reinterpret_cast<const uint8_t*>(image_bgr.data), image_size};
|
|
|
|
auto image_metadata = new_metadata();
|
|
image_metadata->set_string("dtype", "uint8");
|
|
image_metadata->set_string("encoding", "bgr8");
|
|
image_metadata->set_string("layout", "HWC");
|
|
image_metadata->set_list_int("shape", make_shape(height, width, 3));
|
|
image_metadata->set_float("fx", calib.fx);
|
|
image_metadata->set_float("fy", calib.fy);
|
|
image_metadata->set_float("cx", calib.cx);
|
|
image_metadata->set_float("cy", calib.cy);
|
|
image_metadata->set_list_float("distortion", make_distortion(calib.distortion));
|
|
image_metadata->set_int("frame_id", current_frame_id);
|
|
image_metadata->set_int("timestamp_ns", timestamp_ns);
|
|
|
|
auto image_result = send_output_with_metadata(
|
|
dora_node.send_output, "image_bgr", image_slice, std::move(image_metadata));
|
|
if (!std::string(image_result.error).empty()) {
|
|
std::cerr << "Error sending image_bgr: " << std::string(image_result.error) << std::endl;
|
|
}
|
|
|
|
// --- point_cloud ---
|
|
int pc_width = point_cloud_mat.getWidth();
|
|
int pc_height = point_cloud_mat.getHeight();
|
|
|
|
std::vector<float> pc_buf;
|
|
const uint8_t* pc_data_ptr = nullptr;
|
|
size_t pc_data_size = 0;
|
|
auto pc_metadata = new_metadata();
|
|
pc_metadata->set_string("dtype", "float32");
|
|
pc_metadata->set_string("channels", "XYZRGBA");
|
|
pc_metadata->set_string("units", "mm");
|
|
pc_metadata->set_float("fx", calib.fx);
|
|
pc_metadata->set_float("fy", calib.fy);
|
|
pc_metadata->set_float("cx", calib.cx);
|
|
pc_metadata->set_float("cy", calib.cy);
|
|
pc_metadata->set_list_float("distortion", make_distortion(calib.distortion));
|
|
|
|
if (pc_mode == PcMode::HALF) {
|
|
int hw = pc_width / 2, hh = pc_height / 2;
|
|
pc_buf.resize(static_cast<size_t>(hh * hw * 4));
|
|
const sl::float4* src = point_cloud_mat.getPtr<sl::float4>();
|
|
for (int y = 0; y < hh; y++) {
|
|
for (int x = 0; x < hw; x++) {
|
|
const sl::float4& pt = src[(y * 2) * pc_width + (x * 2)];
|
|
float* dst = pc_buf.data() + (y * hw + x) * 4;
|
|
dst[0] = pt.x; dst[1] = pt.y; dst[2] = pt.z; dst[3] = pt.w;
|
|
}
|
|
}
|
|
pc_metadata->set_string("layout", "HWC");
|
|
pc_metadata->set_list_int("shape", make_shape(hh, hw, 4));
|
|
pc_data_ptr = reinterpret_cast<const uint8_t*>(pc_buf.data());
|
|
pc_data_size = pc_buf.size() * sizeof(float);
|
|
} else if (pc_mode == PcMode::ROI) {
|
|
const auto& objs = detector->objects();
|
|
for (const auto& obj : objs) {
|
|
int x1 = std::max(0, static_cast<int>(obj.bounding_box_2d[0]));
|
|
int y1 = std::max(0, static_cast<int>(obj.bounding_box_2d[1]));
|
|
int x2 = std::min(pc_width - 1, static_cast<int>(obj.bounding_box_2d[4]));
|
|
int y2 = std::min(pc_height - 1, static_cast<int>(obj.bounding_box_2d[5]));
|
|
for (int y = y1; y <= y2; y++) {
|
|
for (int x = x1; x <= x2; x++) {
|
|
sl::float4 pt;
|
|
point_cloud_mat.getValue(x, y, &pt);
|
|
if (!std::isnan(pt.x) && !std::isinf(pt.x)) {
|
|
pc_buf.push_back(pt.x); pc_buf.push_back(pt.y);
|
|
pc_buf.push_back(pt.z); pc_buf.push_back(pt.w);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
int N = static_cast<int>(pc_buf.size() / 4);
|
|
pc_metadata->set_string("layout", "NC");
|
|
pc_metadata->set_string("mode", "roi");
|
|
pc_metadata->set_list_int("shape", make_vector_int({N, 4}));
|
|
pc_data_ptr = reinterpret_cast<const uint8_t*>(pc_buf.data());
|
|
pc_data_size = pc_buf.size() * sizeof(float);
|
|
} else {
|
|
// FULL — send raw sl::Mat pointer directly (no copy)
|
|
pc_data_size = static_cast<size_t>(pc_width * pc_height * sizeof(sl::float4));
|
|
pc_data_ptr = reinterpret_cast<const uint8_t*>(point_cloud_mat.getPtr<sl::float4>());
|
|
pc_metadata->set_string("layout", "HWC");
|
|
pc_metadata->set_list_int("shape", make_shape(pc_height, pc_width, 4));
|
|
}
|
|
|
|
rust::Slice<const uint8_t> pc_slice{pc_data_ptr, pc_data_size};
|
|
auto pc_result = send_output_with_metadata(
|
|
dora_node.send_output, "point_cloud", pc_slice, std::move(pc_metadata));
|
|
if (!std::string(pc_result.error).empty()) {
|
|
std::cerr << "Error sending point_cloud: " << std::string(pc_result.error) << std::endl;
|
|
}
|
|
|
|
// --- camera_info ---
|
|
float camera_data[4] = {calib.fx, calib.fy, calib.cx, calib.cy};
|
|
rust::Slice<const uint8_t> cam_slice{
|
|
reinterpret_cast<const uint8_t*>(camera_data), sizeof(camera_data)};
|
|
|
|
auto cam_metadata = new_metadata();
|
|
cam_metadata->set_string("dtype", "float32");
|
|
cam_metadata->set_string("layout", "C");
|
|
cam_metadata->set_list_int("shape", make_vector_int({4}));
|
|
cam_metadata->set_list_float("distortion", make_distortion(calib.distortion));
|
|
|
|
auto cam_result = send_output_with_metadata(
|
|
dora_node.send_output, "camera_info", cam_slice, std::move(cam_metadata));
|
|
if (!std::string(cam_result.error).empty()) {
|
|
std::cerr << "Error sending camera_info: " << std::string(cam_result.error) << std::endl;
|
|
}
|
|
|
|
// --- objects ---
|
|
const auto& objects = detector->objects();
|
|
auto obj_buffer = serialize_objects(objects);
|
|
int num_objects = static_cast<int>(objects.size());
|
|
|
|
const size_t obj_size = obj_buffer.size() * sizeof(float);
|
|
rust::Slice<const uint8_t> obj_slice{
|
|
reinterpret_cast<const uint8_t*>(obj_buffer.data()), obj_size};
|
|
|
|
auto obj_metadata = new_metadata();
|
|
obj_metadata->set_string("dtype", "float32");
|
|
obj_metadata->set_string("layout", "RC");
|
|
obj_metadata->set_string("units", "mm");
|
|
obj_metadata->set_list_int("shape", make_vector_int({num_objects, FIELDS_PER_OBJECT}));
|
|
obj_metadata->set_string("fields",
|
|
"label_id,confidence,x,y,z,"
|
|
"bb2d_x0,bb2d_y0,bb2d_x1,bb2d_y1,bb2d_x2,bb2d_y2,bb2d_x3,bb2d_y3,"
|
|
"bb3d_x0,bb3d_y0,bb3d_z0,bb3d_x1,bb3d_y1,bb3d_z1,"
|
|
"bb3d_x2,bb3d_y2,bb3d_z2,bb3d_x3,bb3d_y3,bb3d_z3,"
|
|
"bb3d_x4,bb3d_y4,bb3d_z4,bb3d_x5,bb3d_y5,bb3d_z5,"
|
|
"bb3d_x6,bb3d_y6,bb3d_z6,bb3d_x7,bb3d_y7,bb3d_z7");
|
|
obj_metadata->set_string("label_names", detector->label_names());
|
|
obj_metadata->set_string("engine", engine_name);
|
|
obj_metadata->set_int("frame_id", current_frame_id);
|
|
obj_metadata->set_int("timestamp_ns", timestamp_ns);
|
|
obj_metadata->set_int("image_width", width);
|
|
obj_metadata->set_int("image_height", height);
|
|
|
|
auto obj_result = send_output_with_metadata(
|
|
dora_node.send_output, "objects", obj_slice, std::move(obj_metadata));
|
|
if (!std::string(obj_result.error).empty()) {
|
|
std::cerr << "Error sending objects: " << std::string(obj_result.error) << std::endl;
|
|
}
|
|
}
|
|
|
|
return 0;
|
|
}
|