std::vector NLClassifier::BuildResults()

in tensorflow_lite_support/cc/task/text/nlclassifier/nl_classifier.cc [141:181]

• 37 lines of code
• 14 McCabe index (conditional complexity)

std::vector<Category> NLClassifier::BuildResults(const TfLiteTensor* scores,
                                                 const TfLiteTensor* labels) {
  bool use_index_as_labels = (labels_vector_ == nullptr) && (labels == nullptr);
  // Some models output scores with transposed shape [1, categories]
  int categories =
      scores->dims->size == 2 ? scores->dims->data[1] : scores->dims->data[0];

  std::vector<Category> predictions;
  predictions.reserve(categories);

  bool should_dequantize = scores->type == kTfLiteUInt8 ||
                           scores->type == kTfLiteInt8 ||
                           scores->type == kTfLiteInt16;
  for (int index = 0; index < categories; index++) {
    std::string label;
    if (use_index_as_labels) {
      label = std::to_string(index);
    } else if (labels_vector_ == nullptr) {
      if (labels->type == kTfLiteString) {
        label = GetStringAtIndex(labels, index);
      } else if (labels->type == kTfLiteInt32) {
        label = std::to_string(GetTensorData<int>(labels)[index]);
      }
    } else {
      label = (*labels_vector_)[index];
    }
    if (should_dequantize) {
      predictions.push_back(Category(label, Dequantize(*scores, index)));
    } else if (scores->type == kTfLiteBool) {
      predictions.push_back(
          Category(label, GetTensorData<bool>(scores)[index] ? 1.0 : 0.0));
    } else {
      predictions.push_back(
          Category(label, scores->type == kTfLiteFloat32
                              ? GetTensorData<float>(scores)[index]
                              : GetTensorData<double>(scores)[index]));
    }
  }

  return predictions;
}