Program Listing for File yolo_detector.cpp#
↰ Return to documentation for file (Src/yolo_detector.cpp)
module;
#include "kataglyphis_export.h"
#include <algorithm>
#include <array>
#include <atomic>
#include <cmath>
#include <cstring>
#include <expected>
#include <filesystem>
#include <functional>
#include <mutex>
#include <numeric>
#include <shared_mutex>
#include <span>
#include <utility>
module kataglyphis.yolo_detector;
import kataglyphis.onnx_inference;
import kataglyphis.gstreamer_pipeline;
import kataglyphis.project_config;
namespace kataglyphis::detection {
namespace {
constexpr std::array<const char *, 80> COCO_CLASSES = { "person",
"bicycle",
"car",
"motorcycle",
"airplane",
"bus",
"train",
"truck",
"boat",
"traffic light",
"fire hydrant",
"stop sign",
"parking meter",
"bench",
"bird",
"cat",
"dog",
"horse",
"sheep",
"cow",
"elephant",
"bear",
"zebra",
"giraffe",
"backpack",
"umbrella",
"handbag",
"tie",
"suitcase",
"frisbee",
"skis",
"snowboard",
"sports ball",
"kite",
"baseball bat",
"baseball glove",
"skateboard",
"surfboard",
"tennis racket",
"bottle",
"wine glass",
"cup",
"fork",
"knife",
"spoon",
"bowl",
"banana",
"apple",
"sandwich",
"orange",
"broccoli",
"carrot",
"hot dog",
"pizza",
"donut",
"cake",
"chair",
"couch",
"potted plant",
"bed",
"dining table",
"toilet",
"tv",
"laptop",
"mouse",
"remote",
"keyboard",
"cell phone",
"microwave",
"oven",
"toaster",
"sink",
"refrigerator",
"book",
"clock",
"vase",
"scissors",
"teddy bear",
"hair drier",
"toothbrush", };
auto copy_buffer_to_float_data(const gstreamer::BufferInfo &buffer) -> std::expected<std::vector<float>, OnnxError>
{
if (buffer.data.empty()) { return std::unexpected(OnnxError::InputAllocationFailed); }
if (buffer.size % sizeof(float) != 0) { return std::unexpected(OnnxError::InvalidInputShape); }
if (buffer.data.size() < buffer.size) { return std::unexpected(OnnxError::InputAllocationFailed); }
const std::size_t num_floats = buffer.size / sizeof(float);
std::vector<float> float_data(num_floats);
std::memcpy(float_data.data(), buffer.data.data(), buffer.size);
return float_data;
}
auto build_camera_detection_pipeline_description(const std::string &device,
std::uint32_t width,
std::uint32_t height,
std::uint32_t fps) -> std::string
{
return "v4l2src device=" + device + " ! "
"video/x-raw,width="
+ std::to_string(width) + ",height=" + std::to_string(height) + ",framerate="
+ std::to_string(fps)
+ "/1 ! "
"videoconvert ! video/x-raw,format=RGB ! "
"appsink name=detector_sink emit-signals=true sync=false";
}
auto build_video_detection_pipeline_description(const std::string &video_source,
std::uint32_t width,
std::uint32_t height) -> std::string
{
std::string pipeline_desc;
if (video_source.contains("://")) {
pipeline_desc = "uridecodebin uri=" + video_source + " ! ";
} else {
pipeline_desc = "filesrc location=" + video_source + " ! decodebin ! ";
}
pipeline_desc +=
"videoconvert ! videoscale ! "
"video/x-raw,format=RGB,width="
+ std::to_string(width) + ",height=" + std::to_string(height)
+ " ! appsink name=detector_sink emit-signals=true sync=false";
return pipeline_desc;
}
// Frame pixels per model-input pixel, by which every decoded box is scaled.
struct Scale
{
float x;
float y;
};
auto labelled_box(float x, float y, float width, float height, float confidence, int class_id) -> BoundingBox
{
BoundingBox box;
box.x = x;
box.y = y;
box.width = width;
box.height = height;
box.confidence = confidence;
box.class_id = class_id;
box.class_name = YoloDetector::get_coco_class_name(class_id);
return box;
}
auto decode_center_rows(std::span<const float> data,
std::size_t rows,
std::size_t values,
const YoloConfig &config,
Scale scale) -> std::vector<BoundingBox>
{
std::vector<BoundingBox> boxes;
// A row is cx, cy, w, h, objectness, then class scores; anything shorter is not a detection layout.
if (values < 5) { return boxes; }
// Class scores beyond the row would be read from the next detection.
const std::size_t class_count = std::min(static_cast<std::size_t>(std::max(config.num_classes, 0)), values - 5);
for (std::size_t i = 0; i < rows && (i + 1) * values <= data.size(); ++i) {
const auto row = data.subspan(i * values, values);
const float cx = row[0];
const float cy = row[1];
const float w = row[2];
const float h = row[3];
const float obj_conf = row[4];
// NaN fails every comparison, so the thresholds test for acceptance, not rejection.
if (!(obj_conf >= config.confidence_threshold)) { continue; }
if (!std::isfinite(cx) || !std::isfinite(cy) || !std::isfinite(w) || !std::isfinite(h)) { continue; }
int best_class = 0;
float final_conf = obj_conf;
if (class_count > 0) {
float best_class_conf = 0.0F;
for (std::size_t c = 0; c < class_count; ++c) {
if (row[5 + c] > best_class_conf) {
best_class_conf = row[5 + c];
best_class = static_cast<int>(c);
}
}
final_conf = obj_conf * best_class_conf;
}
if (!(final_conf >= config.confidence_threshold) || !std::isfinite(final_conf)) { continue; }
boxes.push_back(labelled_box((cx - (w / 2.0F)) * scale.x,
(cy - (h / 2.0F)) * scale.y,
w * scale.x,
h * scale.y,
final_conf,
best_class));
}
return boxes;
}
// An end-to-end row worth a box: a finite score over the threshold, a known class, ordered finite corners.
auto corner_row_is_usable(std::span<const float> row, const YoloConfig &config) -> bool
{
if (!(row[4] >= config.confidence_threshold) || !std::isfinite(row[4])) { return false; }
if (!(row[5] >= 0.0F) || !(row[5] < static_cast<float>(config.num_classes))) { return false; }
const bool finite = std::ranges::all_of(row.first(4), [](float v) -> bool { return std::isfinite(v); });
return finite && row[2] >= row[0] && row[3] >= row[1];
}
auto decode_corner_rows(std::span<const float> data,
std::size_t rows,
std::size_t values,
const YoloConfig &config,
Scale scale) -> std::vector<BoundingBox>
{
std::vector<BoundingBox> boxes;
if (values < 6) { return boxes; }
for (std::size_t i = 0; i < rows && (i + 1) * values <= data.size(); ++i) {
const auto row = data.subspan(i * values, values);
if (!corner_row_is_usable(row, config)) { continue; }
boxes.push_back(labelled_box(row[0] * scale.x,
row[1] * scale.y,
(row[2] - row[0]) * scale.x,
(row[3] - row[1]) * scale.y,
row[4],
static_cast<int>(row[5])));
}
return boxes;
}
}// namespace
struct YoloDetector::Impl
{
// detect reads config while initialize rewrites it; the engine guards only its own session.
mutable std::shared_mutex mutex;
inference::OnnxInferenceEngine engine;
YoloConfig config;
bool initialized{ false };
};
YoloDetector::YoloDetector() : impl_(std::make_unique<Impl>()) {}
YoloDetector::~YoloDetector() = default;
YoloDetector::YoloDetector(YoloDetector &&) noexcept = default;
auto YoloDetector::operator=(YoloDetector &&) noexcept -> YoloDetector & = default;
auto YoloDetector::initialize(const YoloConfig &config) -> std::expected<void, OnnxError>
{
const std::unique_lock lock(impl_->mutex);
impl_->initialized = false;
impl_->config = config;
auto session_config = inference::create_default_session_config(config.model_path);
auto result = impl_->engine.initialize(session_config);
if (!result) { return std::unexpected(result.error()); }
impl_->initialized = true;
return {};
}
auto YoloDetector::is_initialized() const -> bool
{
const std::shared_lock lock(impl_->mutex);
return impl_->initialized;
}
auto YoloDetector::detect(std::span<const float> image_data, std::uint32_t width, std::uint32_t height)
-> std::expected<DetectionResult, OnnxError>
{
const std::shared_lock lock(impl_->mutex);
if (!impl_->initialized) { return std::unexpected(OnnxError::SessionNotInitialized); }
inference::TensorShape input_shape;
input_shape.dimensions = { 1, 3, impl_->config.input_height, impl_->config.input_width };
auto result = impl_->engine.run_inference(image_data, input_shape, "images");
if (!result) { return std::unexpected(result.error()); }
return post_process(result.value(), width, height);
}
auto YoloDetector::detect_from_gstreamer(gstreamer::GStreamerPipeline &pipeline, std::uint32_t timeout_ms)
-> std::expected<DetectionResult, OnnxError>
{
if (!is_initialized()) { return std::unexpected(OnnxError::SessionNotInitialized); }
auto buffer_result = pipeline.pull_sample(timeout_ms);
if (!buffer_result) { return std::unexpected(OnnxError::InputAllocationFailed); }
const auto &buffer_info = buffer_result.value();
auto float_data = copy_buffer_to_float_data(buffer_info);
if (!float_data) { return std::unexpected(float_data.error()); }
return detect(std::span<const float>(float_data->data(), float_data->size()),
buffer_info.metadata.width,
buffer_info.metadata.height);
}
auto YoloDetector::post_process(const inference::InferenceResult &raw_output,
std::uint32_t original_width,
std::uint32_t original_height) -> DetectionResult
{
DetectionResult result;
result.original_width = original_width;
result.original_height = original_height;
result.inference_time_ms = raw_output.inference_time_ms;
if (raw_output.outputs.empty()) { return result; }
const auto &output = raw_output.outputs[0];
const auto &dims = output.shape.dimensions;
if (dims.empty()) { return result; }
std::size_t num_detections = dims[0];
std::size_t values_per_detection = dims.size() > 1 ? dims[1] : dims[0];
if (dims.size() == 3) {
num_detections = dims[1];
values_per_detection = dims[2];
}
const std::span<const float> data(output.data.data(), output.data.size());
const Scale scale{ .x=static_cast<float>(original_width) / static_cast<float>(impl_->config.input_width),
.y=static_cast<float>(original_height) / static_cast<float>(impl_->config.input_height), };
auto all_boxes = impl_->config.output_layout == YoloOutputLayout::CornersScoreClass
? decode_corner_rows(data, num_detections, values_per_detection, impl_->config, scale)
: decode_center_rows(data, num_detections, values_per_detection, impl_->config, scale);
apply_nms(all_boxes, impl_->config.nms_threshold);
result.boxes = std::move(all_boxes);
return result;
}
auto YoloDetector::apply_nms(std::vector<BoundingBox> &boxes, float nms_threshold) -> void
{
std::ranges::sort(
boxes, [](const BoundingBox &a, const BoundingBox &b) -> bool { return a.confidence > b.confidence; });
std::vector<bool> suppressed(boxes.size(), false);
for (std::size_t i = 0; i < boxes.size(); ++i) {
if (suppressed[i]) { continue; }
for (std::size_t j = i + 1; j < boxes.size(); ++j) {
if (suppressed[j]) { continue; }
if (boxes[i].class_id == boxes[j].class_id) {
float iou = calculate_iou(boxes[i], boxes[j]);
if (iou > nms_threshold) { suppressed[j] = true; }
}
}
}
std::vector<BoundingBox> kept;
for (std::size_t i = 0; i < boxes.size(); ++i) {
if (!suppressed[i]) { kept.push_back(std::move(boxes[i])); }
}
boxes = std::move(kept);
}
auto YoloDetector::calculate_iou(const BoundingBox &a, const BoundingBox &b) -> float
{
float x1 = std::max(a.x, b.x);
float y1 = std::max(a.y, b.y);
float x2 = std::min(a.x + a.width, b.x + b.width);
float y2 = std::min(a.y + a.height, b.y + b.height);
if (x2 <= x1 || y2 <= y1) { return 0.0F; }
float intersection = (x2 - x1) * (y2 - y1);
float area_a = a.width * a.height;
float area_b = b.width * b.height;
float union_area = area_a + area_b - intersection;
return intersection / union_area;
}
auto YoloDetector::get_coco_class_name(int class_id) -> std::string_view
{
if (class_id >= 0 && std::cmp_less(class_id, COCO_CLASSES.size())) { return COCO_CLASSES[class_id]; }
return "unknown";
}
struct VideoDetectorPipeline::Impl
{
YoloDetector detector;
std::function<void(const DetectionResult &, const gstreamer::BufferInfo &)> detection_callback;
std::function<void(const gstreamer::BufferInfo &)> frame_callback;
std::atomic<bool> running{ false };
VideoDetectionConfig config;
// Last, so it is destroyed first: its streaming thread calls into every member above.
gstreamer::GStreamerPipeline pipeline;
void on_frame_received(const gstreamer::BufferInfo &buffer)
{
if (frame_callback) { frame_callback(buffer); }
if (detector.is_initialized()) {
auto float_data = copy_buffer_to_float_data(buffer);
if (!float_data) { return; }
auto detection = detector.detect(std::span<const float>(float_data->data(), float_data->size()),
buffer.metadata.width,
buffer.metadata.height);
if (detection && detection_callback) { detection_callback(detection.value(), buffer); }
}
}
};
VideoDetectorPipeline::VideoDetectorPipeline() : impl_(std::make_unique<Impl>()) {}
VideoDetectorPipeline::~VideoDetectorPipeline() = default;
VideoDetectorPipeline::VideoDetectorPipeline(VideoDetectorPipeline &&) noexcept = default;
auto VideoDetectorPipeline::operator=(VideoDetectorPipeline &&) noexcept -> VideoDetectorPipeline & = default;
auto VideoDetectorPipeline::initialize(const VideoDetectionConfig &config) -> std::expected<void, OnnxError>
{
impl_->config = config;
auto init_result = gstreamer::GStreamerPipeline::initialize_gstreamer();
if (!init_result) { return std::unexpected(OnnxError::SessionCreationFailed); }
auto detector_result = impl_->detector.initialize(config.yolo_config);
if (!detector_result) { return std::unexpected(detector_result.error()); }
auto pipeline_result = impl_->pipeline.create_pipeline(config.gstreamer_config);
if (!pipeline_result) { return std::unexpected(OnnxError::SessionCreationFailed); }
// The Impl, not `this`: the factories return the pipeline by move, which leaves `this` behind.
impl_->pipeline.set_buffer_callback(
[impl = impl_.get()](const gstreamer::BufferInfo &buffer) -> void { impl->on_frame_received(buffer); });
return {};
}
auto VideoDetectorPipeline::start() -> std::expected<void, gstreamer::GStreamerError>
{
auto result = impl_->pipeline.start();
if (result) { impl_->running.store(true); }
return result;
}
auto VideoDetectorPipeline::stop() -> std::expected<void, gstreamer::GStreamerError>
{
impl_->running.store(false);
return impl_->pipeline.stop();
}
auto VideoDetectorPipeline::pause() -> std::expected<void, gstreamer::GStreamerError>
{ return impl_->pipeline.pause(); }
auto VideoDetectorPipeline::set_detection_callback(
std::function<void(const DetectionResult &, const gstreamer::BufferInfo &)> callback) -> void
{ impl_->detection_callback = std::move(callback); }
auto VideoDetectorPipeline::set_frame_callback(std::function<void(const gstreamer::BufferInfo &)> callback) -> void
{ impl_->frame_callback = std::move(callback); }
auto VideoDetectorPipeline::is_running() const -> bool { return impl_->running.load(); }
auto create_camera_detection_pipeline(const std::string &device,
const std::filesystem::path &model_path,
std::uint32_t width,
std::uint32_t height,
std::uint32_t fps) -> std::expected<VideoDetectorPipeline, OnnxError>
{
VideoDetectorPipeline pipeline;
VideoDetectionConfig config;
config.yolo_config.model_path = model_path;
config.yolo_config.input_width = width;
config.yolo_config.input_height = height;
config.gstreamer_config.pipeline_description =
build_camera_detection_pipeline_description(device, width, height, fps);
config.gstreamer_config.enable_tensor_meta = true;
auto result = pipeline.initialize(config);
if (!result) { return std::unexpected(result.error()); }
return pipeline;
}
auto create_video_detection_pipeline(const std::string &video_source,
const std::filesystem::path &model_path,
std::uint32_t width,
std::uint32_t height) -> std::expected<VideoDetectorPipeline, OnnxError>
{
VideoDetectorPipeline pipeline;
VideoDetectionConfig config;
config.yolo_config.model_path = model_path;
config.yolo_config.input_width = width;
config.yolo_config.input_height = height;
config.gstreamer_config.pipeline_description =
build_video_detection_pipeline_description(video_source, width, height);
config.gstreamer_config.enable_tensor_meta = true;
auto result = pipeline.initialize(config);
if (!result) { return std::unexpected(result.error()); }
return pipeline;
}
}// namespace kataglyphis::detection