ML_Model.h

Go to the documentation of this file.

/*
 * PROPRIETARY INFORMATION.  This software is proprietary to
 * Side Effects Software Inc., and is not to be reproduced,
 * transmitted, or disclosed in any way without written permission.
 *
 * COMMENTS: Wrapper for the ONNX inference engine
 *
 */

#pragma once

#include "ML_API.h"
#include <UT/UT_Array.h>
#include <UT/UT_NonCopyable.h>
#include <UT/UT_SharedPtr.h>
#include <UT/UT_StringHolder.h>

class UT_WorkBuffer;

class ML_API ML_Model
{
public:
    ML_Model();
    ~ML_Model();
    UT_NON_COPYABLE(ML_Model)

    class SessionInfo; // Our environment
    using Shape = UT_Array<int64>;

    /// Initializer for the ML_Model class
    /// \param model_filepath contains the path to the ONNX model
    /// \param usecuda will trigger if we add the cuda provider or not.
    /// \param errors is meant to hold any error strings that may be generated from
    /// an error occurring during initialization
    bool initializeModel(const UT_StringRef &model_filepath,
                         bool usecuda,
                         UT_WorkBuffer &errors,
                         UT_WorkBuffer &warnings);

    bool run(const UT_Array<UT_Array<float>> &inputs,
             const UT_Array<Shape> &input_shapes,
             UT_Array<UT_Array<float>> &outputs,
             const UT_Array<Shape> &output_shapes,
             UT_WorkBuffer &error_message);

    void getNames(UT_StringArray &input_names,
                  UT_StringArray &output_names) const;

    void getShapes(UT_Array<Shape> &input_shapes,
                   UT_Array<Shape> &output_shapes) const;

    /// Gets the product of all non-dynamic axes of a tensor shape.
    /// Places a bool in a variable to determine if dynamic axes were found
    /// Any dimensions being zero will return 0
    /// \returns 1 if all axes are dynamic
    static exint tensorElementsSize(const UT_Array<exint> &tensor_dimensions,
                                    bool &has_dynamic_axes);

    /// This is a function for acquiring the shape of a tensor from the parameters.
    /// \param tensor_shape the array to fill with the shape.
    static bool mat3ToShape(Shape &tensor_shape, const UT_Matrix3D &shape_vector);

    /// Parses the output data for nodes
    /// \param maxtuplesize -1 for unlimited
    static bool parseOutputData(const UT_StringHolder &srcpattern, int maxtuplesize,
                                UT_StringArray &names, UT_IntArray &sizes, UT_WorkBuffer &error);

private:
    /// Places the information about the model into a UT_WorkBuffer
    void info(UT_WorkBuffer &model_info) const;

    /// Places the string representing the shape for input "input_index" into a workbuffer
    void inputShapeString(int input_index, UT_WorkBuffer &the_string) const;

    /// Places the string representing the shape for output "output_index" into a workbuffer
    void outputShapeString(int output_index, UT_WorkBuffer &the_string) const;

    /// Checks to see if the size and shape of the inputs and outputs
    /// are compatible with the model and should be able to run
    bool sizeAndShapeErrorChecking(const UT_Array<UT_Array<float>> &inputs,
                                   const UT_Array<Shape> &specified_input_shapes,
                                   UT_Array<UT_Array<float>> &outputs,
                                   const UT_Array<Shape> &specified_output_shapes,
                                   UT_WorkBuffer &error_message);

    UT_SharedPtr<SessionInfo> myEnv;

};