From 35cfa8a9064bb1e4ff0aba491b4c82b9a844d6e3 Mon Sep 17 00:00:00 2001 From: uvos Date: Fri, 5 Apr 2024 11:44:05 +0200 Subject: [PATCH] fix yolo network occasinally preditcting a match out side of image bounds --- yolo.cpp | 21 ++++++++++++++++++++- yolo.h | 17 +++++++---------- 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/yolo.cpp b/yolo.cpp index 6218abf..74d9f2a 100644 --- a/yolo.cpp +++ b/yolo.cpp @@ -181,14 +181,33 @@ std::vector Yolo::runInference(const cv::Mat &input) result.className = classes[result.class_id].first; result.priority = classes[result.class_id].second; + clampBox(boxes[idx], input.size()); result.box = boxes[idx]; - detections.push_back(result); } return detections; } + +void Yolo::clampBox(cv::Rect& box, const cv::Size& size) +{ + if(box.x < 0) + { + box.width += box.x; + box.x = 0; + } + if(box.y < 0) + { + box.height += box.y; + box.y = 0; + } + if(box.x+box.width > size.width) + box.width = size.width - box.x; + if(box.y+box.height > size.height) + box.height = size.height - box.y; +} + void Yolo::loadClasses(const std::string& classesStr) { std::vector candidateClasses = tokenizeBinaryIgnore(classesStr, '\n', '"', '\\'); diff --git a/yolo.h b/yolo.h index af62509..8241f9a 100644 --- a/yolo.h +++ b/yolo.h @@ -27,19 +27,16 @@ private: static constexpr float modelScoreThreshold = 0.45; static constexpr float modelNMSThreshold = 0.50; + std::string modelPath; + std::vector> classes; + cv::Size2f modelShape; + bool letterBoxForSquare = true; + cv::dnn::Net net; + void loadClasses(const std::string& classes); void loadOnnxNetwork(const std::filesystem::path& path); cv::Mat formatToSquare(const cv::Mat &source); - - std::string modelPath; - - std::vector> classes; - - cv::Size2f modelShape; - - bool letterBoxForSquare = true; - - cv::dnn::Net net; + static void clampBox(cv::Rect& box, const cv::Size& size); public: Yolo(const std::filesystem::path &onnxModelPath = "", const cv::Size& modelInputShape = {640, 480},