Yolo-Detection/yolo+ByteTrack/yoloTest/main.cpp

241 lines
6.8 KiB
C++
Raw Permalink Normal View History

2024-10-11 20:43:44 +08:00
#include "yolo.h"
2024-10-23 15:30:46 +08:00
#include "kalmanFilter.h"
2024-10-11 20:43:44 +08:00
#include <iostream>
#include <opencv2/opencv.hpp>
#include <vector>
#include <time.h>
2024-10-23 15:30:46 +08:00
#include "BYTETracker.h"
2024-10-11 20:43:44 +08:00
using namespace std;
using namespace cv;
using namespace cv::dnn;
int main()
{
2024-10-23 15:30:46 +08:00
string img_path = "D:/VS/yoloTest/images/detect.mp4";
string model_path = "D:/VS/yoloTest/models/yolov5s.onnx";
2024-10-11 20:43:44 +08:00
// <20><><EFBFBD><EFBFBD>Yolov5ģ<35>ͺ͸<CDBA><CDB8><EFBFBD><EFBFBD><EFBFBD>
Yolov5 yolo;
Net net;
2024-10-23 15:30:46 +08:00
BYTETracker tracker(30, 30); // <20><><EFBFBD><EFBFBD>Sort<72><74><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
2024-10-11 20:43:44 +08:00
// <20><>ȡYoloģ<6F><C4A3>
if (!yolo.readModel(net, model_path, true)) {
cout << "<EFBFBD>޷<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Yoloģ<EFBFBD><EFBFBD>" << endl;
return -1;
}
// <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ɫ<EFBFBD><C9AB><EFBFBD><EFBFBD>ÿ<EFBFBD><C3BF><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ӧһ<D3A6><D2BB><EFBFBD><EFBFBD>ɫ
vector<Scalar> colors;
srand(42); // <20>̶<EFBFBD><CCB6><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ӣ<EFBFBD>ȷ<EFBFBD><C8B7>ÿ<EFBFBD><C3BF><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ɫһ<C9AB><D2BB>
for (int i = 0; i < 80; i++) { // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>80<38><30><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
int b = rand() % 256;
int g = rand() % 256;
int r = rand() % 256;
colors.push_back(Scalar(b, g, r));
}
// <20><>ȡ<EFBFBD><C8A1>Ƶ
VideoCapture cap(img_path);
if (!cap.isOpened()) {
cout << "<EFBFBD>޷<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ƶ<EFBFBD>ļ<EFBFBD>" << endl;
return -1;
}
VideoWriter writer;
2024-10-23 15:30:46 +08:00
writer = VideoWriter("D:/VS/yoloTest/images/test.mp4", CAP_OPENCV_MJPEG, 30, Size(1920, 1080), true);
2024-10-11 20:43:44 +08:00
clock_t start, end; //<2F><>ʱ<EFBFBD><CAB1>
start = clock();
int num = 0;
Mat frame;
int frame_counter = 0;
int detect_interval = 3; // ÿ<><C3BF><EFBFBD><EFBFBD><EFBFBD><EFBFBD>֡<EFBFBD><D6A1><EFBFBD>µ<EFBFBD><C2B5><EFBFBD>һ<EFBFBD><D2BB>YOLO<4C><4F><EFBFBD>м<EFBFBD><D0BC><EFBFBD>
vector<Rect> detection_rects;
2024-10-23 15:30:46 +08:00
vector<Object> objects;
2024-10-11 20:43:44 +08:00
while (cap.read(frame)) {
if (frame.empty()) {
cout << "<EFBFBD><EFBFBD>Ƶ<EFBFBD>ѽ<EFBFBD><EFBFBD><EFBFBD>" << endl;
break;
}
vector<Output> result;
// ÿ<><C3BF> detect_interval ֡<><D6A1><EFBFBD><EFBFBD>һ<EFBFBD><D2BB> YOLO <20><><EFBFBD><EFBFBD>
if (frame_counter % detect_interval == 0) {
detection_rects.clear(); // <20><><EFBFBD><EFBFBD>֮ǰ<D6AE>ļ<EFBFBD><C4BC><EFBFBD><EFBFBD><EFBFBD>
2024-10-23 15:30:46 +08:00
objects.clear();
2024-10-11 20:43:44 +08:00
// <20><><EFBFBD><EFBFBD>Ŀ<EFBFBD><C4BF><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
if (yolo.Detect(frame, net, result)) {
for (int i = 0; i < result.size(); i++) {
int x = result[i].box.x;
int y = result[i].box.y;
int w = result[i].box.width;
int h = result[i].box.height;
Rect rect(x, y, w, h);
2024-10-23 15:30:46 +08:00
Object object;
object.label = result[i].id;
object.prob = result[i].confidence;
object.rect = rect;
objects.push_back(object);
2024-10-11 20:43:44 +08:00
}
}
// <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ݸ<EFBFBD> SORT <20><><EFBFBD>г<EFBFBD>ʼ<EFBFBD><CABC><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
2024-10-23 15:30:46 +08:00
vector<STrack> object_trackers = tracker.update(objects);
2024-10-11 20:43:44 +08:00
// <20><><EFBFBD>Ƹ<EFBFBD><C6B8>ٽ<EFBFBD><D9BD><EFBFBD>
2024-10-23 15:30:46 +08:00
vector<vector<float>> trackers;
for (int i = 0; i < object_trackers.size(); i++) {
trackers.push_back(object_trackers[i].tlwh);
}
2024-10-11 20:43:44 +08:00
for (int i = 0; i < trackers.size(); i++) {
2024-10-23 15:30:46 +08:00
Rect rect(trackers[i][0], trackers[i][1], trackers[i][2], trackers[i][3]);
2024-10-11 20:43:44 +08:00
// <20><>ȡ<EFBFBD><C8A1><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ŷ<EFBFBD>
// <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFA3AC><EFBFBD>Ǽ<EFBFBD><C7BC><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ŷ<EFBFBD>û<EFBFBD>б
int class_id = 0; // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD> ID Ϊ 0<><30><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ʾ<EFBFBD><CABE>
float confidence = 0.8f; // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ŷ<EFBFBD>Ϊ 0.8<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ʾ<EFBFBD><EFBFBD>
// ʹ<><CAB9><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ɫ<EFBFBD><C9AB><EFBFBD>л<EFBFBD>ȡ<EFBFBD><C8A1>ɫ
Scalar color = colors[class_id % colors.size()]; // ʹ<><CAB9><EFBFBD><EFBFBD>ɫ<EFBFBD><C9AB><EFBFBD>е<EFBFBD><D0B5><EFBFBD>ɫ
// <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ǩ<EFBFBD><C7A9><EFBFBD>ݣ<EFBFBD><DDA3><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD> + <20><><EFBFBD>Ŷ<EFBFBD>
string label = "box";
// <20><><EFBFBD>ƾ<EFBFBD><C6BE>ο<EFBFBD><CEBF><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ǩ
rectangle(frame, rect, color, 2); // <20><><EFBFBD>ƾ<EFBFBD><C6BE>ο<EFBFBD>
putText(frame, label, Point(rect.x, rect.y), FONT_HERSHEY_SIMPLEX, 1, color, 2); // <20><>ʾ<EFBFBD><CABE><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ƺ<EFBFBD><C6BA><EFBFBD><EFBFBD>Ŷ<EFBFBD>
}
}
else {
// <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ǽ<EFBFBD><C7BC><EFBFBD>֡<EFBFBD><D6A1>ʹ<EFBFBD><CAB9> SORT <20><><EFBFBD><EFBFBD>Ŀ<EFBFBD><C4BF>Ԥ<EFBFBD><D4A4><EFBFBD>͸<EFBFBD><CDB8><EFBFBD>
2024-10-23 15:30:46 +08:00
vector<STrack> object_trackers = tracker.update(objects);
// <20><><EFBFBD>Ƹ<EFBFBD><C6B8>ٽ<EFBFBD><D9BD><EFBFBD>
vector<vector<float>> trackers;
for (int i = 0; i < object_trackers.size(); i++) {
trackers.push_back(object_trackers[i].tlwh);
}
2024-10-11 20:43:44 +08:00
// <20><><EFBFBD>Ƹ<EFBFBD><C6B8>ٽ<EFBFBD><D9BD><EFBFBD>
for (int i = 0; i < trackers.size(); i++) {
2024-10-23 15:30:46 +08:00
Rect rect(trackers[i][0], trackers[i][1], trackers[i][2], trackers[i][3]);
2024-10-11 20:43:44 +08:00
// <20><>ȡ<EFBFBD><C8A1><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ŷ<EFBFBD>
// <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFA3AC><EFBFBD>Ǽ<EFBFBD><C7BC><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ŷ<EFBFBD>û<EFBFBD>б
int class_id = 0; // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD> ID Ϊ 0<><30><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ʾ<EFBFBD><CABE>
float confidence = 0.8f; // <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ŷ<EFBFBD>Ϊ 0.8<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ʾ<EFBFBD><EFBFBD>
// ʹ<><CAB9><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ɫ<EFBFBD><C9AB><EFBFBD>л<EFBFBD>ȡ<EFBFBD><C8A1>ɫ
Scalar color = colors[class_id % colors.size()]; // ʹ<><CAB9><EFBFBD><EFBFBD>ɫ<EFBFBD><C9AB><EFBFBD>е<EFBFBD><D0B5><EFBFBD>ɫ
// <20><><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ǩ<EFBFBD><C7A9><EFBFBD>ݣ<EFBFBD><DDA3><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD> + <20><><EFBFBD>Ŷ<EFBFBD>
string label = "box";
// <20><><EFBFBD>ƾ<EFBFBD><C6BE>ο<EFBFBD><CEBF><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ǩ
rectangle(frame, rect, color, 2); // <20><><EFBFBD>ƾ<EFBFBD><C6BE>ο<EFBFBD>
putText(frame, label, Point(rect.x, rect.y), FONT_HERSHEY_SIMPLEX, 1, color, 2); // <20><>ʾ<EFBFBD><CABE><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ƺ<EFBFBD><C6BA><EFBFBD><EFBFBD>Ŷ<EFBFBD>
}
}
// <20><>ʾ<EFBFBD><CABE><EFBFBD><EFBFBD><EFBFBD><EFBFBD>֡
imshow("img", frame);
writer << frame;
if (waitKey(10) == 27) {
break; // <20>˳<EFBFBD>ѭ<EFBFBD><D1AD>
}
frame_counter++;
num++;
}
end = clock();
cout << "time = " << double(end - start) / CLOCKS_PER_SEC << "s" << endl; // <20><><EFBFBD><EFBFBD>ʱ<EFBFBD><EFBFBD><E4A3A8>λ<EFBFBD><CEBB>s<EFBFBD><73>
cout << "frame num: " << num << endl;
cout << "Speed = " << double(end - start) * 1000 / CLOCKS_PER_SEC / num << "ms" << endl; // <20><><EFBFBD><EFBFBD>ʱ<EFBFBD><EFBFBD><E4A3A8>λ<EFBFBD><CEBB>ms<6D><73>
cap.release();
destroyAllWindows();
2024-10-23 15:30:46 +08:00
writer.release();
2024-10-11 20:43:44 +08:00
return 0;
}
2024-10-15 21:31:44 +08:00
/*
#include "yolo.h"
#include <iostream>
2024-10-23 15:30:46 +08:00
#include<opencv2//opencv.hpp>
#include<math.h>
2024-10-15 21:31:44 +08:00
#include <time.h>
2024-10-23 15:30:46 +08:00
2024-10-15 21:31:44 +08:00
using namespace std;
using namespace cv;
2024-10-23 15:30:46 +08:00
using namespace dnn;
2024-10-15 21:31:44 +08:00
int main()
{
2024-10-23 15:30:46 +08:00
string img_path = "D:/VS/yoloTest/images/2.mp4";
string model = "D:/VS/yoloTest/models/yolov5s.onnx";
Yolov5 test;
Net net;
if (test.readModel(net, model, false)) {
cout << "read net ok!" << endl;
}
else {
return -1;
}
//<2F><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ɫ
vector<Scalar> color;
srand(time(0));
for (int i = 0; i < 80; i++) {
int b = rand() % 256;
int g = rand() % 256;
int r = rand() % 256;
color.push_back(Scalar(b, g, r));
}
vector<Output> result;
VideoCapture cap(img_path); //<2F><>ȡͼ<C8A1><CDBC>/<2F><>Ƶ
VideoWriter writer;
writer = VideoWriter("D:/VS/yoloTest/images/test.mp4", CAP_OPENCV_MJPEG, 20, Size(2560, 1440), true); //<2F><><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ϊ<EFBFBD><CEAA>Ƶ
clock_t start, end; //<2F><>ʱ<EFBFBD><CAB1>
start = clock();
int num = 0;
Mat frame;
while (cap.read(frame)) {
vector<Output> result;
if (test.Detect(frame, net, result)) {
test.drawPred(frame, result, color);
}
imshow("img", frame);
writer << frame;
if (waitKey(10) == 27) {
break; // <20>˳<EFBFBD>ѭ<EFBFBD><D1AD>
}
num++;
}
end = clock();
cout << "time = " << double(end - start) / CLOCKS_PER_SEC << "s" << endl; //<2F><><EFBFBD><EFBFBD>ʱ<EFBFBD><EFBFBD><E4A3A8>λ<EFBFBD><CEBB><EFBFBD><EFBFBD><EFBFBD><EFBFBD>
cout << "frame num: " << num << endl;
cout << "Speed = " << double(end - start)*1000 / CLOCKS_PER_SEC/num << "ms" << endl; //<2F><><EFBFBD><EFBFBD>ʱ<EFBFBD><EFBFBD><E4A3A8>λ<EFBFBD><CEBB>m<EFBFBD><6D><EFBFBD><EFBFBD>
cap.release();
destroyAllWindows();
system("pause");
return 0;
2024-10-15 21:31:44 +08:00
}
*/