| 55 | } |
| 56 | |
| 57 | NeuralNetworkDetector::ResultsVec NeuralNetworkDetector::detect(const std::vector<std::string>& labels) const |
| 58 | { |
| 59 | if (!session_) { |
| 60 | LogError << "OrtSession not loaded"; |
| 61 | return { }; |
| 62 | } |
| 63 | |
| 64 | // batch_size, channel, height, width |
| 65 | // for yolov8, input_shape is { 1, 3, 640, 640 } |
| 66 | const auto input_shape = session_->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape(); |
| 67 | if (input_shape.size() != 4) { |
| 68 | LogError << "Input shape is not 4" << VAR(input_shape); |
| 69 | return { }; |
| 70 | } |
| 71 | |
| 72 | cv::Mat image = image_with_roi(); |
| 73 | cv::Size raw_roi_size(image.cols, image.rows); |
| 74 | cv::Size input_image_size(static_cast<int>(input_shape[3]), static_cast<int>(input_shape[2])); |
| 75 | cv::resize(image, image, input_image_size, 0, 0, cv::INTER_AREA); |
| 76 | std::vector<float> input = image_to_tensor(image); |
| 77 | |
| 78 | Ort::Value input_tensor = |
| 79 | Ort::Value::CreateTensor<float>(memory_info_, input.data(), input.size(), input_shape.data(), input_shape.size()); |
| 80 | |
| 81 | Ort::AllocatorWithDefaultOptions allocator; |
| 82 | const std::string in_0 = session_->GetInputNameAllocated(0, allocator).get(); |
| 83 | const std::string out_0 = session_->GetOutputNameAllocated(0, allocator).get(); |
| 84 | const std::vector input_names { in_0.c_str() }; |
| 85 | const std::vector output_names { out_0.c_str() }; |
| 86 | |
| 87 | Ort::RunOptions run_options; |
| 88 | auto output_tensor = |
| 89 | session_->Run(run_options, input_names.data(), &input_tensor, input_names.size(), output_names.data(), output_names.size()); |
| 90 | |
| 91 | const float* raw_output = output_tensor[0].GetTensorData<float>(); |
| 92 | // output_shape is { 1, 5, 8400 } |
| 93 | std::vector<int64_t> output_shape = output_tensor[0].GetTensorTypeAndShapeInfo().GetShape(); |
| 94 | |
| 95 | // yolov8 的 onnx 输出和前面的 v5, v7 等似乎不太一样，目前网上 yolov8 的 demo 较少，文档也没找到 |
| 96 | // 这里的输出解析是我跟着数据推测的： |
| 97 | // center_x0, center_x1, ..... center_x8399 |
| 98 | // center_y0, center_y1, ..... center_y8399 |
| 99 | // w0, w1, ..... w8399 |
| 100 | // h0, h1, ..... h8399 |
| 101 | // cls1: conf0, conf1, ..... conf8399 |
| 102 | // cls2: conf0, conf1, ..... conf8399 |
| 103 | // cls3: conf0, conf1, ..... conf8399 |
| 104 | // ...... |
| 105 | std::vector<std::vector<float>> output(output_shape[1]); |
| 106 | for (int64_t i = 0; i < output_shape[1]; i++) { |
| 107 | output[i] = std::vector<float>(raw_output + i * output_shape[2], raw_output + (i + 1) * output_shape[2]); |
| 108 | } |
| 109 | |
| 110 | ResultsVec raw_results; |
| 111 | const size_t output_size = output.back().size(); |
| 112 | double width_ratio = 1.0 * raw_roi_size.width / input_image_size.width; |
| 113 | double height_ratio = 1.0 * raw_roi_size.height / input_image_size.height; |
| 114 | |