-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmain.cpp
More file actions
executable file
·67 lines (59 loc) · 1.86 KB
/
Copy pathmain.cpp
File metadata and controls
executable file
·67 lines (59 loc) · 1.86 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
#include<iostream>
#include "manager.hpp"
#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/opencv.hpp>
#include <vector>
#include <chrono>
#include <map>
#include <cmath>
#include <time.h>
using namespace cv;
using namespace std;
//Note: the input model files can be in *.onnx format;
const char* yolo_engine = "../resources/models/yolov5s.engine";
const char* sort_engine = "../resources/models/deepsort.engine";
float conf_threshold = 0.4;
// output configuration
const char* output_video_name = "result.mp4";
const int fps = 30;
string get_input_video_path(int argc, char *argv[]) {
string video_source = "";
if (argc < 2) {
std::cout << "[trace] insufficient arguments provided, exiting..." << std::endl;
return video_source;
} else {
std::cout << "[trace] input offline video source: " << video_source << std::endl;
return argv[1];
}
}
int run_main(int argc, char *argv[]) {
const string video_source = get_input_video_path(argc, argv);
map<int, vector<int>> personstate;
map<int, int> classidmap;
bool is_first = true;
Trtyolosort yosort(yolo_engine, sort_engine);
VideoCapture capture;
cv::Mat frame;
frame = capture.open(video_source);
if (!capture.isOpened()) {
std::cerr << "[error] can not open target file: " << video_source << std::endl;
return -1;
}
capture.read(frame);
const int frame_width = frame.cols;
const int frame_height = frame.rows;
std::cout << "[trace] original video dimension: width=" << frame_width << " height=" << frame_height << std::endl;
// configurable part for final output video
const int fourcc = cv::VideoWriter::fourcc('M','J','P','G');
std::vector<DetectBox> det;
while (capture.read(frame)) {
yosort.TrtDetect(frame, conf_threshold, det);
}
capture.release();
return 0;
}
int main (int argc, char *argv[]) {
const int result = run_main(argc, argv);
return result;
}