wangrong
2025-03-19 291afb9c81aae1d15d9e3edcea618d0a2e1fc5a6
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
 
 
#include "../trt_yolov8_detector.h"
#include "../trt_yolov8_pose_detector.h"
#include "../trt_yolov8_seg_detector.h"
#include "../trt_yolov8_classifier.h"
 
int main(int argc, char** argv) {
    /* run command:
     * ./trt_yolov8_wts_2_engine [-det/-seg/-pose/-cls] [.wts] [.engine] [n/s/m/l/x/n2/s2/m2/l2/x2/n6/s6/m6/l6/x6] 
    */
    
    if (argc != 5) {
        std::cerr << "arguments not right!" << std::endl;
        std::cerr << "./trt_yolov8_wts_2_engine [-det/-seg/-pose/-cls] [.wts] [.engine] [n/s/m/l/x/n2/s2/m2/l2/x2/n6/s6/m6/l6/x6]" << std::endl;
    }
    
    std::string task_type = std::string(argv[1]);
    std::string wts_name = std::string(argv[2]);
    std::string engine_name = std::string(argv[3]);
    std::string sub_type = std::string(argv[4]);
 
    if (task_type == "-det") {
        trt_yolov8::trt_yolov8_detector detector;
        detector.wts_2_engine(wts_name, engine_name, sub_type);
    }
    else if (task_type == "-seg") {
        trt_yolov8::trt_yolov8_seg_detector detector;
        detector.wts_2_engine(wts_name, engine_name, sub_type);
    }
    else if (task_type == "-pose") {
        trt_yolov8::trt_yolov8_pose_detector detector;
        detector.wts_2_engine(wts_name, engine_name, sub_type);
    }
    else if (task_type == "-cls") {
        trt_yolov8::trt_yolov8_classifier classifier;
        classifier.wts_2_engine(wts_name, engine_name, sub_type);
    }
    else {
        std::cerr << "arguments not right!" << std::endl;
        std::cerr << "./trt_yolov8_wts_2_engine [-det/-seg/-pose/-cls] [.wts] [.engine] [n/s/m/l/x/n2/s2/m2/l2/x2/n6/s6/m6/l6/x6]" << std::endl;
    }
 
    return 0;
}