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
#include "vehicle_scanner.h"
 
 
namespace trt_vehicle {
    VehicleScanner::VehicleScanner(const std::string& scanModelPath, float scoreThres, float iouThres) {
        scannerModel = new DetectModel(scanModelPath, scoreThres, iouThres);
    }
    
    VehicleScanner::~VehicleScanner() {
        if (scannerModel) {
            delete scannerModel;
        }
    }
    
    void VehicleScanner::detect(std::vector<cv::Mat>& img_datas, std::vector<std::vector<trt_vehicle::ObjBox>>& parts) {
        parts.clear();
        // batch size == 1 for vehicle det model
        // scan 1 by 1
        for (int i = 0; i < img_datas.size(); i++) {
            std::vector<std::vector<ObjBox>> outBoxes;
            scannerModel->predictPadding({img_datas[i]}, outBoxes, 128);
 
            auto& outBox = outBoxes[0];
            std::vector<trt_vehicle::ObjBox> part_list;
 
            if (outBox.size() == 0) {
                parts.push_back(part_list);
                continue;
            }
 
            for (auto box: outBox) {
                // attach vehicle label field
                box.label = partLabels[box.class_];
                part_list.push_back(box);
            }
            
            parts.push_back(part_list);
        }
    } 
}