| 46 | } |
| 47 | |
| 48 | NeuralNetworkClassifier::Result NeuralNetworkClassifier::classify() const |
| 49 | { |
| 50 | if (!session_) { |
| 51 | LogError << "OrtSession not loaded"; |
| 52 | return { }; |
| 53 | } |
| 54 | // batch_size, channel, height, width |
| 55 | // for yolov8, input_shape is { 1, 3, 640, 640 } |
| 56 | const auto input_shape = session_->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape(); |
| 57 | if (input_shape.size() != 4) { |
| 58 | LogError << "Input shape is not 4" << VAR(input_shape); |
| 59 | return { }; |
| 60 | } |
| 61 | |
| 62 | cv::Mat image = image_with_roi(); |
| 63 | cv::Size raw_roi_size(image.cols, image.rows); |
| 64 | cv::Size input_image_size(static_cast<int>(input_shape[3]), static_cast<int>(input_shape[2])); |
| 65 | cv::resize(image, image, input_image_size, 0, 0, cv::INTER_AREA); |
| 66 | std::vector<float> input = image_to_tensor(image); |
| 67 | |
| 68 | Ort::Value input_tensor = |
| 69 | Ort::Value::CreateTensor<float>(memory_info_, input.data(), input.size(), input_shape.data(), input_shape.size()); |
| 70 | |
| 71 | Ort::AllocatorWithDefaultOptions allocator; |
| 72 | const std::string in_0 = session_->GetInputNameAllocated(0, allocator).get(); |
| 73 | const std::string out_0 = session_->GetOutputNameAllocated(0, allocator).get(); |
| 74 | const std::vector input_names { in_0.c_str() }; |
| 75 | const std::vector output_names { out_0.c_str() }; |
| 76 | |
| 77 | Ort::RunOptions run_options; |
| 78 | auto output_tensor = |
| 79 | session_->Run(run_options, input_names.data(), &input_tensor, input_names.size(), output_names.data(), output_names.size()); |
| 80 | |
| 81 | const float* raw_output = output_tensor[0].GetTensorData<float>(); |
| 82 | std::vector<float> output(raw_output, raw_output + output_tensor[0].GetTensorTypeAndShapeInfo().GetElementCount()); |
| 83 | |
| 84 | Result res; |
| 85 | res.raw = std::move(output); |
| 86 | res.probs = softmax(res.raw); |
| 87 | res.cls_index = std::max_element(res.probs.begin(), res.probs.end()) - res.probs.begin(); |
| 88 | res.score = res.probs[res.cls_index]; |
| 89 | res.label = res.cls_index < param_.labels.size() ? param_.labels[res.cls_index] : std::format("Unknown_{}", res.cls_index); |
| 90 | res.box = roi_; |
| 91 | |
| 92 | if (debug_draw_) { |
| 93 | auto draw = draw_result(res); |
| 94 | handle_draw(draw); |
| 95 | } |
| 96 | |
| 97 | return res; |
| 98 | } |
| 99 | |
| 100 | void NeuralNetworkClassifier::add_results(ResultsVec results, const std::vector<int>& expected) |
| 101 | { |