initial commit

wangrong

2025-01-03 2c4b8838f89bebca97615c5f0bac2c58e0349771

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; 
}