Method onTrigger

extensions/tensorflow/TFExtractTopLabels.cpp:52–124

}

void TFExtractTopLabels::onTrigger(const std::shared_ptr<core::ProcessContext> &context,
                                   const std::shared_ptr<core::ProcessSession> &session) {
  auto flow_file = session->get();

  if (!flow_file) {
    return;
  }

  try {

    // Read labels
    std::string tf_type;
    flow_file->getAttribute("tf.type", tf_type);
    std::shared_ptr<std::vector<std::string>> labels;

    {
      std::lock_guard<std::mutex> guard(labels_mtx_);

      if (tf_type == "labels") {
        logger_->log_info("Reading new labels...");
        auto new_labels = std::make_shared<std::vector<std::string>>();
        LabelsReadCallback cb(new_labels);
        session->read(flow_file, &cb);
        labels_ = new_labels;
        logger_->log_info("Read %d new labels", labels_->size());
        session->remove(flow_file);
        return;
      }

      labels = labels_;
    }

    // Read input tensor from flow file
    auto input_tensor_proto = std::make_shared<tensorflow::TensorProto>();
    TensorReadCallback tensor_cb(input_tensor_proto);
    session->read(flow_file, &tensor_cb);

    tensorflow::Tensor input;
    input.FromProto(*input_tensor_proto);
    auto input_flat = input.flat<float>();

    std::vector<std::pair<uint64_t, float>> scores;

    for (int i = 0; i < input_flat.size(); i++) {
      scores.emplace_back(std::make_pair(i, input_flat(i)));
    }

    std::sort(scores.begin(), scores.end(), [](const std::pair<uint64_t, float> &a,
                                               const std::pair<uint64_t, float> &b) {
      return a.second > b.second;
    });

    for (int i = 0; i < 5 && i < scores.size(); i++) {
      if (!labels || scores[i].first > labels->size()) {
        logger_->log_error("Label index is out of range (are the correct labels loaded?); routing to retry...");
        session->transfer(flow_file, Retry);
        return;
      }