vtkONNXInference.h

Go to the documentation of this file.

// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#ifndef vtkONNXInference_h
#define vtkONNXInference_h
 
#include "vtkFiltersONNXModule.h" // For export macro
#include "vtkPassInputTypeAlgorithm.h"
 
#include "vtkDataObject.h" // for AttributeTypes
 
#include <memory> // For std::unique_ptr
#include <vector> // For std::vector
 
VTK_ABI_NAMESPACE_BEGIN
struct vtkONNXInferenceInternals;
namespace Ort
{
class AllocatorWithDefaultOptions;
class Value;
}
 
class VTKFILTERSONNX_EXPORT vtkONNXInference : public vtkPassInputTypeAlgorithm
{
public:
  vtkTypeMacro(vtkONNXInference, vtkPassInputTypeAlgorithm);
  void PrintSelf(ostream& os, vtkIndent indent) override;
 
  static vtkONNXInference* New();
 
 
  void SetModelFile(const std::string& file);
  vtkGetMacro(ModelFile, std::string);
 
 
  void SetTimeStepValues(const std::vector<double>& times);
 
  void SetTimeStepValue(vtkIdType idx, double timeStepValue);
 
  void SetNumberOfTimeStepValues(vtkIdType nb);
 
  void ClearTimeStepValues();
 
  vtkSetMacro(TimeStepIndex, int);
 
  vtkGetMacro(TimeStepIndex, int);
 
 
  void SetInputParameters(const std::vector<float>& params);
 
  void SetInputParameter(vtkIdType idx, float InputParameter);
 
  void SetNumberOfInputParameters(vtkIdType nb);
 
  vtkGetMacro(InputSize, int);
 
  void ClearInputParameters();
 
 
  vtkSetMacro(OutputDimension, int);
  vtkGetMacro(OutputDimension, int);
 
 
  vtkSetMacro(ArrayAssociation, int);
  vtkGetMacro(ArrayAssociation, int);
 
protected:
  vtkONNXInference();
  ~vtkONNXInference() override = default;
 
  int RequestInformation(vtkInformation*, vtkInformationVector**, vtkInformationVector*) override;
 
  int RequestData(vtkInformation*, vtkInformationVector**, vtkInformationVector*) override;
 
  int ExecuteData(vtkDataObject* input, vtkDataObject* output, double timevalue);
 
private:
  vtkONNXInference(const vtkONNXInference&) = delete;
  void operator=(const vtkONNXInference&) = delete;
 
  void InitializeSession();
 
  bool ShouldGenerateTimeSteps();
 
  std::vector<Ort::Value> RunModel(Ort::Value& inputTensor);
 
  // Input related parameters
  std::string ModelFile;
  int64_t InputSize = 0;
  std::vector<float> InputParameters;
  std::vector<double> TimeStepValues;
  int TimeStepIndex = -1;
 
  // Output related parameters
  int OutputDimension = 1;
 
  int ArrayAssociation = vtkDataObject::CELL;
 
  bool Initialized = false;
  std::unique_ptr<vtkONNXInferenceInternals> Internals;
};
VTK_ABI_NAMESPACE_END
 
#endif // vtkONNXInference_h