mirror of https://github.com/opencv/opencv.git
Merge pull request #24898 from Abdurrahheem:ash/yolo_ducumentation
Documentation for Yolo usage in Opencv #24898 This PR introduces documentation for the usage of yolo detection model family in open CV. This is not to be merge before #24691, as the sample will need to be changed. ### Pull Request Readiness Checklist See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request - [x] I agree to contribute to the project under Apache 2 License. - [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV - [x] The PR is proposed to the proper branch - [x] There is a reference to the original bug report and related work - [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable Patch to opencv_extra has the same branch name. - [x] The feature is well documented and sample code can be built with the project CMakepull/24945/head
Abduragim Shtanchaev
1 year ago
committed by
GitHub
parent
8e43c8f200
commit
372b36c1d3
2 changed files with 573 additions and 25 deletions
@ -0,0 +1,370 @@ |
|||||||
|
/**
|
||||||
|
* @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& test_name |
||||||
|
); |
||||||
|
|
||||||
|
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. }" |
||||||
|
"{ 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& test_name) |
||||||
|
{ |
||||||
|
// Retrieve
|
||||||
|
std::vector<int> classIds; |
||||||
|
std::vector<float> confidences; |
||||||
|
std::vector<Rect2d> boxes; |
||||||
|
|
||||||
|
if (test_name == "yolov8") |
||||||
|
{ |
||||||
|
cv::transposeND(outs[0], {0, 2, 1}, outs[0]); |
||||||
|
} |
||||||
|
|
||||||
|
if (test_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, 84}); |
||||||
|
} |
||||||
|
|
||||||
|
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 = (test_name == "yolov8" || test_name == "yolonas") ? 1.0f : preds.at<float>(i, 4) ; |
||||||
|
if (obj_conf < conf_threshold) |
||||||
|
continue; |
||||||
|
|
||||||
|
Mat scores = preds.row(i).colRange((test_name == "yolov8" || test_name == "yolonas") ? 4 : 5, preds.cols); |
||||||
|
double conf; |
||||||
|
Point maxLoc; |
||||||
|
minMaxLoc(scores, 0, &conf, 0, &maxLoc); |
||||||
|
|
||||||
|
conf = (test_name == "yolov8" || test_name == "yolonas") ? 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 (test_name == "yolonas"){ |
||||||
|
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 = findFile(parser.get<String>("model")); |
||||||
|
std::string yolo_model = parser.get<String>("yolo"); |
||||||
|
|
||||||
|
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 != "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; |
||||||
|
|
||||||
|
// 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); |
||||||
|
//![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; |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
Loading…
Reference in new issue