Analyzes the output of the graph to retrieve the highest scores and their positions in the tensor.
| 72 | // Analyzes the output of the graph to retrieve the highest scores and |
| 73 | // their positions in the tensor. |
| 74 | void GetTopLabels(const std::vector<Tensor>& outputs, int how_many_labels, |
| 75 | Tensor* out_indices, Tensor* out_scores) { |
| 76 | const Tensor& unsorted_scores_tensor = outputs[0]; |
| 77 | auto unsorted_scores_flat = unsorted_scores_tensor.flat<float>(); |
| 78 | std::vector<std::pair<int, float>> scores; |
| 79 | scores.reserve(unsorted_scores_flat.size()); |
| 80 | for (int i = 0; i < unsorted_scores_flat.size(); ++i) { |
| 81 | scores.push_back(std::pair<int, float>({i, unsorted_scores_flat(i)})); |
| 82 | } |
| 83 | std::sort(scores.begin(), scores.end(), |
| 84 | [](const std::pair<int, float>& left, |
| 85 | const std::pair<int, float>& right) { |
| 86 | return left.second > right.second; |
| 87 | }); |
| 88 | scores.resize(how_many_labels); |
| 89 | Tensor sorted_indices(tensorflow::DT_INT32, {how_many_labels}); |
| 90 | Tensor sorted_scores(tensorflow::DT_FLOAT, {how_many_labels}); |
| 91 | for (int i = 0; i < scores.size(); ++i) { |
| 92 | sorted_indices.flat<int>()(i) = scores[i].first; |
| 93 | sorted_scores.flat<float>()(i) = scores[i].second; |
| 94 | } |
| 95 | *out_indices = sorted_indices; |
| 96 | *out_scores = sorted_scores; |
| 97 | } |
| 98 | |
| 99 | } // namespace |
| 100 |