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 ClearInputParameters();


  void SetInputShape(const std::vector<int64_t>& shape);
  void SetInputShape(vtkIdType idx, int shapeElement);
  void SetInputShape(vtkIdType nb);
  const std::vector<int64_t>& GetInputShape() const;

  void SetNumberOfInputShapeElements(vtkIdType nb);

  void ClearInputShape();


  void SetInputPermutation(const std::vector<int>& shape);
  const std::vector<int>& GetInputPermutation() const;


  void SetOutputPermutation(const std::vector<int>& permutation);
  const std::vector<int>& GetOutputPermutation() const;


  vtkSetMacro(FieldArrayInput, bool);
  vtkGetMacro(FieldArrayInput, bool);
  vtkBooleanMacro(FieldArrayInput, bool);


  vtkSetMacro(ProcessedFieldArrayName, const std::string&);
  vtkGetMacro(ProcessedFieldArrayName, const std::string&);


  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();

  bool GenerateInputTensorFromParameters(
    std::vector<float>& parameters, Ort::Value& inputTensor, double timeValue);

  bool GenerateInputTensorFromFieldArray(
    Ort::Value& inputTensor, vtkDataSetAttributes* inAttributes);

  std::vector<Ort::Value> RunModel(Ort::Value& inputTensor);

  // Input related parameters
  std::string ModelFile;
  std::vector<int64_t> InputShape = { 0 };
  std::vector<float> InputParameters;
  std::vector<double> TimeStepValues;
  int TimeStepIndex = -1;
  bool FieldArrayInput = false;
  std::string ProcessedFieldArrayName;
  std::vector<int> InputPermutation;

  // Output related parameters
  int OutputDimension = 1;
  std::vector<int> OutputPermutation;

  int ArrayAssociation = vtkDataObject::CELL;
  std::vector<float> InputDataBuffer;

  bool Initialized = false;
  std::unique_ptr<vtkONNXInferenceInternals> Internals;
};
VTK_ABI_NAMESPACE_END

#endif // vtkONNXInference_h