| 9 | } |
| 10 | |
| 11 | void Detector::Initialize(int gpu_id, int width, int height, |
| 12 | std::string name_list_path) { |
| 13 | if (gpu_id >= 0) { |
| 14 | if (gpu_id >= torch::getNumGPUs()) { |
| 15 | std::cout << "No GPU id " << gpu_id << " abailable" << std::endl; |
| 16 | } |
| 17 | device = torch::Device(torch::kCUDA, gpu_id); |
| 18 | } |
| 19 | else { |
| 20 | device = torch::Device(torch::kCPU); |
| 21 | } |
| 22 | name_list = {}; |
| 23 | std::ifstream ifs; |
| 24 | ifs.open(name_list_path, std::ios::in); |
| 25 | if (!ifs.is_open()) |
| 26 | { |
| 27 | std::cout<< "Open "<< name_list_path<<" file failed."; |
| 28 | return; |
| 29 | } |
| 30 | std::string buf = ""; |
| 31 | while (getline(ifs, buf)) |
| 32 | { |
| 33 | name_list.push_back(buf); |
| 34 | } |
| 35 | |
| 36 | |
| 37 | int num_classes = name_list.size(); |
| 38 | this->name_list = name_list; |
| 39 | |
| 40 | this->width = width; |
| 41 | this->height = height; |
| 42 | if (width % 32 || height % 32) { |
| 43 | std::cout << "Width or height is not divisible by 32" << std::endl; |
| 44 | return ; |
| 45 | } |
| 46 | |
| 47 | detector = YoloBody_tiny(3, num_classes); |
| 48 | detector->to(device); |
| 49 | return; |
| 50 | } |
| 51 | |
| 52 | |
| 53 | void Detector::loadPretrained(std::string pretrained_pth) { |