Tensor.h

Go to the documentation of this file.

#pragma once
 
#include <variant>
#include <vector>
#include <cstdint>
#include <string>
#include <cmath>
#include <numeric>
#include <iostream>
#include <sstream>
#include <array>
#include <type_traits>
#include <utility>
 
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
 
#include <tl/expected.hpp>
 
#include <ReduceOps.h>
 
#include "NeuralError.h"
 
namespace NN
{
    using Shape = std::vector<int64_t>;
 
    int64_t shapeNumElements(const Shape &shape);
 
    int64_t shapeToStride(const Shape &shape, uint64_t axis);
 
    std::string shapeToString(const Shape &shape);
 
    enum class TensorType
    {
        Float,  // float32
        Int64,  // int64_t
        UInt8,  // uint8_t
        Invalid // Invalid type, used for error handling
    };
 
    constexpr const char *toString(TensorType type)
    {
        switch (type)
        {
        case TensorType::Float:
            return "float32";
        case TensorType::Int64:
            return "int64";
        case TensorType::UInt8:
            return "uint8";
        default:
            return "Invalid";
        }
    };
 
    // --- Helper to extract std::array traits
 
    template <typename T>
    struct is_std_array : std::false_type
    {
    };
 
    template <typename U, std::size_t N>
    struct is_std_array<std::array<U, N>> : std::true_type
    {
        using value_type = U;
        static constexpr size_t size = N;
    };
 
    template <typename T>
    using decay_t = typename std::decay<T>::type;
 
    template <typename Func, typename InputElem>
    struct map_array_traits
    {
        using return_type = std::decay_t<std::invoke_result_t<Func, InputElem>>;
        static_assert(is_std_array<return_type>::value, "Function must return std::array");
 
        using value_type = typename is_std_array<return_type>::value_type;
        static constexpr size_t size = is_std_array<return_type>::size;
    };
    // --- End heplers
 
    class Tensor
    {
    public:
        using TensorData = std::variant<
            std::vector<float>,
            std::vector<int64_t>,
            std::vector<uint8_t>>;
 
        static tl::expected<Tensor, NeuralError> fromCvMat(const cv::Mat &mat);
 
        static tl::expected<Tensor, NeuralError> fromCvMat(const cv::Mat &mat, const Shape &shape, float scale = 1.0f, std::vector<float> mean = {}, std::vector<float> std = {});
 
        static tl::expected<Tensor, NeuralError> fromCvMats(const std::vector<cv::Mat> &mats);
 
        static tl::expected<Tensor, NeuralError> fromCvMats(const std::vector<cv::Mat> &mats, const Shape &shape, float scale = 1.0f, std::vector<float> mean = {}, std::vector<float> std = {});
 
        template <typename T>
        static tl::expected<Tensor, NeuralError> fromData(const std::vector<T> &data, const Shape &shape)
        {
            return fromData(std::vector<T>(data), shape); // delegates to move overload
        }
 
        template <typename T>
        static tl::expected<Tensor, NeuralError> fromData(std::vector<T> &&data, const Shape &shape)
        {
            if (static_cast<size_t>(shapeNumElements(shape)) != data.size())
            {
                return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "Data size does not match shape"));
            }
 
            Tensor t;
            t.m_shape = shape;
            t.m_data = std::move(data);
            return t;
        }
 
        template <typename T>
        tl::expected<std::vector<T>, NeuralError> toVector() const
        {
            if (std::holds_alternative<std::vector<T>>(m_data))
            {
                return std::get<std::vector<T>>(m_data);
            }
            return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "Tensor does not hold requested data type"));
        }
 
        tl::expected<cv::Mat, NeuralError> toCvMat() const;
 
        const Shape &shape() const { return m_shape; }
 
        std::string toString() const;
 
        tl::expected<void, NeuralError> reshape(const Shape &newShape);
 
        int64_t numElements() const;
 
        bool empty() const
        {
            return m_shape.empty() || shapeNumElements(m_shape) == 0;
        }
 
        TensorType dtype() const
        {
            return std::visit([](const auto &vec) -> TensorType
                              {
                using T = typename std::decay_t<decltype(vec)>::value_type;
                if constexpr (std::is_same_v<T, float>) {
                    return TensorType::Float;
                } else if constexpr (std::is_same_v<T, int64_t>) {
                    return TensorType::Int64;
                } else if constexpr (std::is_same_v<T, uint8_t>) {
                    return TensorType::UInt8;
                } else {
                    return TensorType::Invalid; // Unsupported type
                } }, m_data);
        };
 
        template <typename Op>
        tl::expected<Tensor, NeuralError> reduce(const Op &op, uint64_t axis) const
        {
            if (axis >= m_shape.size())
            {
                return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "Invalid reduction axis"));
            }
 
            return std::visit([&](const auto &inputVec) -> tl::expected<Tensor, NeuralError>
                              {
                using T = typename std::decay_t<decltype(inputVec)>::value_type;
                const int64_t D = m_shape.size();
 
                Shape outShape = m_shape;
                outShape[axis] = 1;
 
                std::vector<T> output(shapeNumElements(outShape));
 
                int64_t innerStride = shapeToStride(m_shape, axis);
                int64_t dimSize = m_shape[axis];
                int64_t outerStride = innerStride * dimSize;
 
                for (int64_t offset = 0; offset < inputVec.size(); offset += outerStride) {
                    for (int64_t i = 0; i < innerStride; ++i) {
                        // Compute flat index
                        int64_t outIdx = (offset / outerStride) * innerStride + i;
                        // Initialize accumulator
                        T acc = op.template initial<T>();
                        for (int64_t d = 0; d < dimSize; ++d) {
                            int64_t idx = offset + d * innerStride + i;
                            op(acc, inputVec[idx]); // Apply operation
                        }
                        output[outIdx] = acc;
                    }
                }
                return Tensor::fromData(std::move(output), outShape); }, m_data);
        }
 
        template <typename Op>
        tl::expected<Tensor, NeuralError> reduceWithIndex(const Op &op, uint64_t axis) const
        {
            if (axis >= m_shape.size())
            {
                return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "Invalid reduction axis"));
            }
 
            return std::visit([&](const auto &inputVec) -> tl::expected<Tensor, NeuralError>
                              {
                using T = typename std::decay_t<decltype(inputVec)>::value_type;
 
                Shape outShape = m_shape;
                outShape[axis] = 1;
 
                std::vector<int64_t> output(shapeNumElements(outShape));
 
                int64_t innerStride = shapeToStride(m_shape, axis);
                int64_t dimSize = m_shape[axis];
                int64_t outerStride = innerStride * dimSize;
 
                for (int64_t offset = 0; offset < inputVec.size(); offset += outerStride) {
                    for (int64_t i = 0; i < innerStride; ++i) {
                        int64_t outIdx = (offset / outerStride) * innerStride + i;
                        auto acc = op.template initial<T>();
                        for (size_t d = 0; d < dimSize; ++d) {
                            size_t idx = offset + d * innerStride + i;
                            op(acc, inputVec[idx], d);
                        }
                        output[outIdx] = static_cast<int64_t>(acc.first);
                    }
                }
 
                return Tensor::fromData(std::move(output), outShape); }, m_data);
        }
 
        template <typename Func>
        tl::expected<Tensor, NeuralError> map(Func &&f) const
        {
            return std::visit([&](const auto &inputVec) -> tl::expected<Tensor, NeuralError>
                              {
                                  using T = typename std::decay_t<decltype(inputVec)>::value_type;
                                  using V = std::decay_t<decltype(f(std::declval<T>()))>;
 
                                  std::vector<V> output;
                                  output.resize(inputVec.size());
 
                                  std::transform(inputVec.begin(), inputVec.end(), output.begin(),
                                                 [&](const T &val) -> V
                                                 {
                                                     return f(val);
                                                 });
 
                                  return Tensor::fromData(std::move(output), m_shape); },
                              m_data);
        }
 
        template <typename Func>
        tl::expected<Tensor, NeuralError> map(Func &&func, int axis) const
        {
            // Sanity check axis
            if (axis < 0 || axis > static_cast<int>(m_shape.size()))
            {
                return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "Invalid axis to insert new dimension."));
            }
 
            return std::visit([&](const auto &input) -> tl::expected<Tensor, NeuralError>
                              {
                using T = typename std::decay_t<decltype(input)>::value_type;
                using Traits = map_array_traits<Func, T>;
                using U = typename Traits::value_type;
                constexpr size_t N = Traits::size;
            
                if (input.empty()) {
                    return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "Tensor data is empty."));
                }
            
                // Compute the original and new shapes
                Shape newShape = m_shape;
                newShape.insert(newShape.begin() + axis, N);
            
                const int64_t oldSize = shapeNumElements(m_shape);
                const int64_t newSize = oldSize * N;
            
                std::vector<U> outData(newSize);
            
                // Precompute strides
                std::vector<int64_t> oldStrides = computeStrides(m_shape);
                std::vector<int64_t> newStrides = computeStrides(newShape);
            
                
                // Main loop: for each index in the old tensor, place mapped values in the new buffer
                for (int64_t idx = 0; idx < oldSize; ++idx) {
                    // Convert flat idx to N-dimensional index
                    std::vector<int64_t> coord = unravelIndex(idx, oldStrides);
                
                    // Apply function to get array<U, N>
                    auto mapped = func(input[idx]);
                
                    // Insert mapped[N] into new buffer at axis
                    coord.insert(coord.begin() + axis, 0);
                    for (int64_t i = 0; i < static_cast<int64_t>(N); ++i) {
                        coord[axis] = i;
                        int64_t flatIdx = ravelIndex(coord, newStrides);
                        outData[flatIdx] = mapped[i];
                    }
                }
            
                return Tensor::fromData(std::move(outData), newShape); }, m_data);
        }
 
        tl::expected<void, NeuralError> squeeze();
 
        tl::expected<void, NeuralError> squeeze(int64_t axis);
 
        tl::expected<void, NeuralError> unsqueeze(int64_t axis);
 
    private:
        TensorData m_data;
        Shape m_shape;
 
        static inline std::vector<int64_t> computeStrides(const Shape &shape)
        {
            std::vector<int64_t> strides(shape.size(), 1);
            for (int i = shape.size() - 2; i >= 0; --i)
                strides[i] = strides[i + 1] * shape[i + 1];
            return strides;
        }
 
        static inline std::vector<int64_t> unravelIndex(int64_t index, const std::vector<int64_t> &strides)
        {
            std::vector<int64_t> coords(strides.size());
            for (size_t i = 0; i < strides.size(); ++i)
            {
                coords[i] = index / strides[i];
                index %= strides[i];
            }
            return coords;
        }
 
        static inline int64_t ravelIndex(const std::vector<int64_t> &coords, const std::vector<int64_t> &strides)
        {
            int64_t index = 0;
            for (size_t i = 0; i < coords.size(); ++i)
                index += coords[i] * strides[i];
            return index;
        }
 
        template <typename T, int CV_TYPE>
        static tl::expected<Tensor, NeuralError> fromCvMatsTyped(const std::vector<cv::Mat> &mats)
        {
            if (mats.empty())
            {
                return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "Input vector of cv::Mat is empty"));
            }
 
            int channels = mats[0].channels();
            int height = mats[0].rows;
            int width = mats[0].cols;
 
            Shape shape = {static_cast<int64_t>(mats.size()), static_cast<int64_t>(channels), static_cast<int64_t>(height), static_cast<int64_t>(width)};
            auto totalSize = shapeNumElements(shape);
 
            std::vector<T> data;
            data.reserve(totalSize);
 
            for (const auto &mat : mats)
            {
                if (mat.depth() != CV_TYPE)
                {
                    return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "All cv::Mat must have the same data type"));
                }
 
                if (mat.channels() != channels || mat.rows != height || mat.cols != width)
                {
                    return tl::unexpected(NeuralError(ErrorCode::InvalidArgument, "All cv::Mat must have the same size and number of channels"));
                }
 
                std::vector<cv::Mat> splitted;
                cv::split(mat, splitted);
 
                for (int c = 0; c < channels; ++c)
                {
                    // data.insert(data.end(), splitted[c].datastart, splitted[c].dataend);
                    const T *channelData = splitted[c].ptr<T>();
                    data.insert(data.end(), channelData, channelData + height * width);
                }
            }
 
            return fromData(std::move(data), shape);
        }
 
        static cv::Mat preprocessCvMat(const cv::Mat &mat, const Shape &shape, float scale)
        {
            cv::Mat tmp;
            const cv::Mat *matPtr = &mat; // points to the mat were currently working with
 
            // if 3d, swap red and blue channels from BGR (opencv-standard) to RGB (anything else)
            if (mat.channels() == 3)
            {
                cv::cvtColor(*matPtr, tmp, cv::COLOR_BGR2RGB);
                matPtr = &tmp;
            }
 
            // resize if necessary
            cv::Size size(shape[shape.size() - 2], shape[shape.size() - 1]);
            if (size != cv::Size(-1,-1) && size != matPtr->size())
            {
                cv::resize(*matPtr, tmp, size, 0, 0, cv::INTER_AREA);
                matPtr = &tmp;
            }
 
            // convert to float and apply scale
            matPtr->convertTo(tmp, CV_32F, scale);
            return tmp;
        }
 
        static cv::Mat preprocessCvMat(const cv::Mat &mat, const Shape &shape, float scale, const std::vector<float> &mean, const std::vector<float> &std)
        {
            cv::Mat tmp = preprocessCvMat(mat, shape, scale); // first do the basic preprocessing
 
            // check if mean and std are provided and have the correct size
            assert(mean.size() == tmp.channels());
            assert(std.size() == tmp.channels());
            // check that std is not zero
            assert(std::all_of(std.begin(), std.end(), [](float s) { return s != 0.0f; }));
 
            // subtract mean from each channel
            int channels = tmp.channels();
            int rows = tmp.rows;
            int cols = tmp.cols * channels;
 
            if (tmp.isContinuous())
            {
                cols *= rows;
                rows = 1;
            }
 
            for (int i = 0; i < rows; ++i)
            {
                float *ptr = tmp.ptr<float>(i);
                for (int j = 0; j < cols; j += channels)
                {
                    for (int c = 0; c < channels; ++c)
                    {
                        ptr[j + c] = (ptr[j + c] - mean[c]) / std[c]; // apply mean and std
                    }
                }
            }
 
            return tmp;
        }
 
        friend class OrtNeuralNet;
    };
}
 
#ifndef TENSOR_DEBUG_PRINT
#ifdef NDEBUG
#define TENSOR_DEBUG_PRINT(tensor) ((void)0);
#else
#include <iostream>
#define TENSOR_DEBUG_PRINT(tensor) (std::cout << (tensor).toString() << std::endl);
#endif
#endif
 
#ifndef SHAPE_DEBUG_PRINT
#ifdef NDEBUG
#define SHAPE_DEBUG_PRINT(shape) ((void)0);
#else
#include <iostream>
#define SHAPE_DEBUG_PRINT(shape) (std::cout << NN::shapeToString(shape) << std::endl);
#endif
#endif
 
#ifndef RETURN_ON_ERROR
#define RETURN_ON_ERROR(expr, retVal)                            \
    do                                                           \
    {                                                            \
        auto _res = (expr);                                      \
        if (!_res)                                               \
        {                                                        \
            std::cerr << "Error: " << _res.error() << std::endl; \
            return (retVal);                                     \
        }                                                        \
    } while (0);
#endif