tensor_cpu.h

#pragma once
#include "linearhostmemory.h"

template<typename, int> class ndarray_ref;

namespace iu
{
template<typename PixelType>
class TensorCpu : public LinearHostMemory<PixelType, 1>
{
public:
    typedef PixelType pixel_type;

    TensorCpu() :
            LinearHostMemory<PixelType, 1>(), samples_(0), channels_(0), height_(0), width_(0)
    {
    }

    virtual ~TensorCpu()
    {
    }

    TensorCpu(const unsigned int N, const unsigned int C, const unsigned int H, const unsigned int W) :
            LinearHostMemory<PixelType, 1>(N * C * H * W), samples_(N), channels_(C), height_(H), width_(W)
    {
    }

    TensorCpu(PixelType* host_data, const unsigned int N, const unsigned int C, const unsigned int H,
            const unsigned int W, bool ext_data_pointer = false) :
            LinearHostMemory<PixelType, 1>(host_data, N * C * H * W, ext_data_pointer), samples_(N), channels_(C), height_(H), width_(
                    W)
    {
    }


    PixelType getPixel(unsigned int n, unsigned int c, unsigned int x, unsigned int y)
    {
        return *this->data(n * (channels_ * height_ * width_) + c * (width_ * height_) + x * width_ + y);
    }

    unsigned int samples() const
    {
        return samples_;
    }

    unsigned int channels() const
    {
        return channels_;
    }

    unsigned int height() const
    {
        return height_;
    }

    unsigned int width() const
    {
        return width_;
    }

  friend std::ostream& operator<<(std::ostream & out,
                                  TensorCpu const& tensor)
  {
    out << "Tensor: height=" << tensor.height() << " width="
        << tensor.width() << " samples="  << tensor.samples() << " channels="
        << tensor.channels() << " onDevice=" << tensor.onDevice();
    return out;
  }

  ndarray_ref<PixelType,4> ref() const;

  TensorCpu(const ndarray_ref<PixelType,4> &x);

private:
    unsigned int samples_;
    unsigned int channels_;
    unsigned int height_;
    unsigned int width_;

private:
  TensorCpu(const TensorCpu&);
  TensorCpu& operator=(const TensorCpu&);
};

} // namespace iu