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