TFHandler.h

Go to the documentation of this file.

///////////////////////////////////////////////////////////////////////
// TFHander
//
// A module to evaluate a tensorflow model in ART. 
//
// \author $Author: Micah Groh
////////////////////////////////////////////////////////////////////////

#ifndef TFHANDLER_H
#define TFHANDLER_H

#include <vector>
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/platform/env.h"

namespace tensorflow
{
  /// Wrapper for Tensorflow which handles construction and prediction
  class TFHandler
  {
  public:
    /// Basic constructor, takes path to model pb
    TFHandler(std::string model, int CPUlimit=1);
    ~TFHandler();
    void Initialize(int CPUlimit);

    std::vector<Tensor> Predict(std::vector<std::pair<std::string,Tensor>> inputs, 
                                std::vector<std::string> outputLabels);

  private:
    Session* fSession;
    std::string fModelPath;
  };
}

#endif