mirror of
https://github.com/azaion/autopilot.git
synced 2026-04-22 10:56:34 +00:00
working yolo default model with opencv
This commit is contained in:
Tuomas Järvinen
@@ -0,0 +1,387 @@
|
||||
/**
|
||||
* @file yolo_detector.cpp
|
||||
* @brief Yolo Object Detection Sample
|
||||
* @author OpenCV team
|
||||
*/
|
||||
|
||||
//![includes]
|
||||
#include <opencv2/dnn.hpp>
|
||||
#include <opencv2/imgproc.hpp>
|
||||
#include <opencv2/imgcodecs.hpp>
|
||||
#include <fstream>
|
||||
#include <sstream>
|
||||
#include "iostream"
|
||||
#include "common.hpp"
|
||||
#include <opencv2/highgui.hpp>
|
||||
//![includes]
|
||||
|
||||
using namespace cv;
|
||||
using namespace cv::dnn;
|
||||
|
||||
void getClasses(std::string classesFile);
|
||||
void drawPrediction(int classId, float conf, int left, int top, int right, int bottom, Mat& frame);
|
||||
void yoloPostProcessing(
|
||||
std::vector<Mat>& outs,
|
||||
std::vector<int>& keep_classIds,
|
||||
std::vector<float>& keep_confidences,
|
||||
std::vector<Rect2d>& keep_boxes,
|
||||
float conf_threshold,
|
||||
float iou_threshold,
|
||||
const std::string& model_name,
|
||||
const int nc
|
||||
);
|
||||
|
||||
std::vector<std::string> classes;
|
||||
|
||||
|
||||
std::string keys =
|
||||
"{ help h | | Print help message. }"
|
||||
"{ device | 0 | camera device number. }"
|
||||
"{ model | onnx/models/yolox_s_inf_decoder.onnx | Default model. }"
|
||||
"{ yolo | yolox | yolo model version. }"
|
||||
"{ input i | | Path to input image or video file. Skip this argument to capture frames from a camera. }"
|
||||
"{ classes | | Optional path to a text file with names of classes to label detected objects. }"
|
||||
"{ nc | 80 | Number of classes. Default is 80 (coming from COCO dataset). }"
|
||||
"{ thr | .5 | Confidence threshold. }"
|
||||
"{ nms | .4 | Non-maximum suppression threshold. }"
|
||||
"{ mean | 0.0 | Normalization constant. }"
|
||||
"{ scale | 1.0 | Preprocess input image by multiplying on a scale factor. }"
|
||||
"{ width | 640 | Preprocess input image by resizing to a specific width. }"
|
||||
"{ height | 640 | Preprocess input image by resizing to a specific height. }"
|
||||
"{ rgb | 1 | Indicate that model works with RGB input images instead BGR ones. }"
|
||||
"{ padvalue | 114.0 | padding value. }"
|
||||
"{ paddingmode | 2 | Choose one of computation backends: "
|
||||
"0: resize to required input size without extra processing, "
|
||||
"1: Image will be cropped after resize, "
|
||||
"2: Resize image to the desired size while preserving the aspect ratio of original image }"
|
||||
"{ backend | 0 | Choose one of computation backends: "
|
||||
"0: automatically (by default), "
|
||||
"1: Halide language (http://halide-lang.org/), "
|
||||
"2: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
||||
"3: OpenCV implementation, "
|
||||
"4: VKCOM, "
|
||||
"5: CUDA }"
|
||||
"{ target | 0 | Choose one of target computation devices: "
|
||||
"0: CPU target (by default), "
|
||||
"1: OpenCL, "
|
||||
"2: OpenCL fp16 (half-float precision), "
|
||||
"3: VPU, "
|
||||
"4: Vulkan, "
|
||||
"6: CUDA, "
|
||||
"7: CUDA fp16 (half-float preprocess) }"
|
||||
"{ async | 0 | Number of asynchronous forwards at the same time. "
|
||||
"Choose 0 for synchronous mode }";
|
||||
|
||||
void getClasses(std::string classesFile)
|
||||
{
|
||||
std::ifstream ifs(classesFile.c_str());
|
||||
if (!ifs.is_open())
|
||||
CV_Error(Error::StsError, "File " + classesFile + " not found");
|
||||
std::string line;
|
||||
while (std::getline(ifs, line))
|
||||
classes.push_back(line);
|
||||
}
|
||||
|
||||
void drawPrediction(int classId, float conf, int left, int top, int right, int bottom, Mat& frame)
|
||||
{
|
||||
rectangle(frame, Point(left, top), Point(right, bottom), Scalar(0, 255, 0));
|
||||
|
||||
std::string label = format("%.2f", conf);
|
||||
if (!classes.empty())
|
||||
{
|
||||
CV_Assert(classId < (int)classes.size());
|
||||
label = classes[classId] + ": " + label;
|
||||
}
|
||||
|
||||
int baseLine;
|
||||
Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
|
||||
|
||||
top = max(top, labelSize.height);
|
||||
rectangle(frame, Point(left, top - labelSize.height),
|
||||
Point(left + labelSize.width, top + baseLine), Scalar::all(255), FILLED);
|
||||
putText(frame, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.5, Scalar());
|
||||
}
|
||||
|
||||
void yoloPostProcessing(
|
||||
std::vector<Mat>& outs,
|
||||
std::vector<int>& keep_classIds,
|
||||
std::vector<float>& keep_confidences,
|
||||
std::vector<Rect2d>& keep_boxes,
|
||||
float conf_threshold,
|
||||
float iou_threshold,
|
||||
const std::string& model_name,
|
||||
const int nc=80)
|
||||
{
|
||||
// Retrieve
|
||||
std::vector<int> classIds;
|
||||
std::vector<float> confidences;
|
||||
std::vector<Rect2d> boxes;
|
||||
|
||||
if (model_name == "yolov8" || model_name == "yolov10" ||
|
||||
model_name == "yolov9")
|
||||
{
|
||||
cv::transposeND(outs[0], {0, 2, 1}, outs[0]);
|
||||
}
|
||||
|
||||
if (model_name == "yolonas")
|
||||
{
|
||||
// outs contains 2 elemets of shape [1, 8400, 80] and [1, 8400, 4]. Concat them to get [1, 8400, 84]
|
||||
Mat concat_out;
|
||||
// squeeze the first dimension
|
||||
outs[0] = outs[0].reshape(1, outs[0].size[1]);
|
||||
outs[1] = outs[1].reshape(1, outs[1].size[1]);
|
||||
cv::hconcat(outs[1], outs[0], concat_out);
|
||||
outs[0] = concat_out;
|
||||
// remove the second element
|
||||
outs.pop_back();
|
||||
// unsqueeze the first dimension
|
||||
outs[0] = outs[0].reshape(0, std::vector<int>{1, 8400, nc + 4});
|
||||
}
|
||||
|
||||
// assert if last dim is 85 or 84
|
||||
CV_CheckEQ(outs[0].dims, 3, "Invalid output shape. The shape should be [1, #anchors, 85 or 84]");
|
||||
CV_CheckEQ((outs[0].size[2] == nc + 5 || outs[0].size[2] == 80 + 4), true, "Invalid output shape: ");
|
||||
|
||||
for (auto preds : outs)
|
||||
{
|
||||
preds = preds.reshape(1, preds.size[1]); // [1, 8400, 85] -> [8400, 85]
|
||||
for (int i = 0; i < preds.rows; ++i)
|
||||
{
|
||||
// filter out non object
|
||||
float obj_conf = (model_name == "yolov8" || model_name == "yolonas" ||
|
||||
model_name == "yolov9" || model_name == "yolov10") ? 1.0f : preds.at<float>(i, 4) ;
|
||||
if (obj_conf < conf_threshold)
|
||||
continue;
|
||||
|
||||
Mat scores = preds.row(i).colRange((model_name == "yolov8" || model_name == "yolonas" || model_name == "yolov9" || model_name == "yolov10") ? 4 : 5, preds.cols);
|
||||
double conf;
|
||||
Point maxLoc;
|
||||
minMaxLoc(scores, 0, &conf, 0, &maxLoc);
|
||||
|
||||
conf = (model_name == "yolov8" || model_name == "yolonas" || model_name == "yolov9" || model_name == "yolov10") ? conf : conf * obj_conf;
|
||||
if (conf < conf_threshold)
|
||||
continue;
|
||||
|
||||
// get bbox coords
|
||||
float* det = preds.ptr<float>(i);
|
||||
double cx = det[0];
|
||||
double cy = det[1];
|
||||
double w = det[2];
|
||||
double h = det[3];
|
||||
|
||||
// [x1, y1, x2, y2]
|
||||
if (model_name == "yolonas" || model_name == "yolov10"){
|
||||
boxes.push_back(Rect2d(cx, cy, w, h));
|
||||
} else {
|
||||
boxes.push_back(Rect2d(cx - 0.5 * w, cy - 0.5 * h,
|
||||
cx + 0.5 * w, cy + 0.5 * h));
|
||||
}
|
||||
classIds.push_back(maxLoc.x);
|
||||
confidences.push_back(static_cast<float>(conf));
|
||||
}
|
||||
}
|
||||
|
||||
// NMS
|
||||
std::vector<int> keep_idx;
|
||||
NMSBoxes(boxes, confidences, conf_threshold, iou_threshold, keep_idx);
|
||||
|
||||
for (auto i : keep_idx)
|
||||
{
|
||||
keep_classIds.push_back(classIds[i]);
|
||||
keep_confidences.push_back(confidences[i]);
|
||||
keep_boxes.push_back(boxes[i]);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @function main
|
||||
* @brief Main function
|
||||
*/
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
CommandLineParser parser(argc, argv, keys);
|
||||
parser.about("Use this script to run object detection deep learning networks using OpenCV.");
|
||||
if (parser.has("help"))
|
||||
{
|
||||
parser.printMessage();
|
||||
return 0;
|
||||
}
|
||||
|
||||
CV_Assert(parser.has("model"));
|
||||
CV_Assert(parser.has("yolo"));
|
||||
// if model is default, use findFile to get the full path otherwise use the given path
|
||||
std::string weightPath = "yolov8n.onnx";
|
||||
std::string yolo_model = "yolov8";
|
||||
int nc = parser.get<int>("nc");
|
||||
|
||||
float confThreshold = parser.get<float>("thr");
|
||||
float nmsThreshold = parser.get<float>("nms");
|
||||
//![preprocess_params]
|
||||
float paddingValue = parser.get<float>("padvalue");
|
||||
bool swapRB = parser.get<bool>("rgb");
|
||||
int inpWidth = parser.get<int>("width");
|
||||
int inpHeight = parser.get<int>("height");
|
||||
Scalar scale = parser.get<float>("scale");
|
||||
Scalar mean = parser.get<Scalar>("mean");
|
||||
ImagePaddingMode paddingMode = static_cast<ImagePaddingMode>(parser.get<int>("paddingmode"));
|
||||
//![preprocess_params]
|
||||
|
||||
// check if yolo model is valid
|
||||
if (yolo_model != "yolov5" && yolo_model != "yolov6"
|
||||
&& yolo_model != "yolov7" && yolo_model != "yolov8"
|
||||
&& yolo_model != "yolov10" && yolo_model !="yolov9"
|
||||
&& yolo_model != "yolox" && yolo_model != "yolonas")
|
||||
CV_Error(Error::StsError, "Invalid yolo model: " + yolo_model);
|
||||
|
||||
// get classes
|
||||
if (parser.has("classes"))
|
||||
{
|
||||
getClasses(findFile(parser.get<String>("classes")));
|
||||
}
|
||||
|
||||
// load model
|
||||
//![read_net]
|
||||
Net net = readNet(weightPath);
|
||||
int backend = parser.get<int>("backend");
|
||||
net.setPreferableBackend(backend);
|
||||
net.setPreferableTarget(parser.get<int>("target"));
|
||||
//![read_net]
|
||||
|
||||
VideoCapture cap;
|
||||
Mat img;
|
||||
bool isImage = false;
|
||||
bool isCamera = false;
|
||||
|
||||
isImage = true;
|
||||
img = imread("bus.png");
|
||||
|
||||
/*
|
||||
// Check if input is given
|
||||
if (parser.has("input"))
|
||||
{
|
||||
String input = parser.get<String>("input");
|
||||
// Check if the input is an image
|
||||
if (input.find(".jpg") != String::npos || input.find(".png") != String::npos)
|
||||
{
|
||||
img = imread(findFile(input));
|
||||
if (img.empty())
|
||||
{
|
||||
CV_Error(Error::StsError, "Cannot read image file: " + input);
|
||||
}
|
||||
isImage = true;
|
||||
}
|
||||
else
|
||||
{
|
||||
cap.open(input);
|
||||
if (!cap.isOpened())
|
||||
{
|
||||
CV_Error(Error::StsError, "Cannot open video " + input);
|
||||
}
|
||||
isCamera = true;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
int cameraIndex = parser.get<int>("device");
|
||||
cap.open(cameraIndex);
|
||||
if (!cap.isOpened())
|
||||
{
|
||||
CV_Error(Error::StsError, cv::format("Cannot open camera #%d", cameraIndex));
|
||||
}
|
||||
isCamera = true;
|
||||
}
|
||||
*/
|
||||
|
||||
// image pre-processing
|
||||
//![preprocess_call]
|
||||
Size size(inpWidth, inpHeight);
|
||||
Image2BlobParams imgParams(
|
||||
scale,
|
||||
size,
|
||||
mean,
|
||||
swapRB,
|
||||
CV_32F,
|
||||
DNN_LAYOUT_NCHW,
|
||||
paddingMode,
|
||||
paddingValue);
|
||||
|
||||
// rescale boxes back to original image
|
||||
Image2BlobParams paramNet;
|
||||
paramNet.scalefactor = scale;
|
||||
paramNet.size = size;
|
||||
paramNet.mean = mean;
|
||||
paramNet.swapRB = swapRB;
|
||||
paramNet.paddingmode = paddingMode;
|
||||
//![preprocess_call]
|
||||
|
||||
//![forward_buffers]
|
||||
std::vector<Mat> outs;
|
||||
std::vector<int> keep_classIds;
|
||||
std::vector<float> keep_confidences;
|
||||
std::vector<Rect2d> keep_boxes;
|
||||
std::vector<Rect> boxes;
|
||||
//![forward_buffers]
|
||||
|
||||
Mat inp;
|
||||
while (waitKey(1) < 0)
|
||||
{
|
||||
|
||||
if (isCamera)
|
||||
cap >> img;
|
||||
if (img.empty())
|
||||
{
|
||||
std::cout << "Empty frame" << std::endl;
|
||||
waitKey();
|
||||
break;
|
||||
}
|
||||
//![preprocess_call_func]
|
||||
inp = blobFromImageWithParams(img, imgParams);
|
||||
//![preprocess_call_func]
|
||||
|
||||
//![forward]
|
||||
net.setInput(inp);
|
||||
net.forward(outs, net.getUnconnectedOutLayersNames());
|
||||
//![forward]
|
||||
|
||||
//![postprocess]
|
||||
yoloPostProcessing(
|
||||
outs, keep_classIds, keep_confidences, keep_boxes,
|
||||
confThreshold, nmsThreshold,
|
||||
yolo_model,
|
||||
nc);
|
||||
//![postprocess]
|
||||
|
||||
// covert Rect2d to Rect
|
||||
//![draw_boxes]
|
||||
for (auto box : keep_boxes)
|
||||
{
|
||||
boxes.push_back(Rect(cvFloor(box.x), cvFloor(box.y), cvFloor(box.width - box.x), cvFloor(box.height - box.y)));
|
||||
}
|
||||
|
||||
paramNet.blobRectsToImageRects(boxes, boxes, img.size());
|
||||
|
||||
for (size_t idx = 0; idx < boxes.size(); ++idx)
|
||||
{
|
||||
Rect box = boxes[idx];
|
||||
drawPrediction(keep_classIds[idx], keep_confidences[idx], box.x, box.y,
|
||||
box.width + box.x, box.height + box.y, img);
|
||||
}
|
||||
|
||||
const std::string kWinName = "Yolo Object Detector";
|
||||
namedWindow(kWinName, WINDOW_NORMAL);
|
||||
imshow(kWinName, img);
|
||||
//![draw_boxes]
|
||||
|
||||
outs.clear();
|
||||
keep_classIds.clear();
|
||||
keep_confidences.clear();
|
||||
keep_boxes.clear();
|
||||
boxes.clear();
|
||||
|
||||
if (isImage)
|
||||
{
|
||||
waitKey();
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user