uvos/SDImagePreprocess
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix yolo network occasinally preditcting a match out side of image bounds

Browse source
This commit is contained in:
uvos 2024-04-05 11:44:05 +02:00
parent a279001151
commit 35cfa8a906
2 changed files with 27 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

21
yolo.cpp
View file

@ -181,14 +181,33 @@ std::vector<Yolo::Detection> 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<std::string> candidateClasses = tokenizeBinaryIgnore(classesStr, '\n', '"', '\\');

17
yolo.h
View file

@ -27,19 +27,16 @@ private:
static constexpr float modelScoreThreshold = 0.45;
static constexpr float modelNMSThreshold = 0.50;
std::string modelPath;
std::vector<std::pair<std::string, int>> 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<std::pair<std::string, int>> 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},