transformer.h

#ifndef SINGA_IO_TRANSFORMER_H_
#define SINGA_IO_TRANSFORMER_H_

#include <vector>
#include <string>
#include "singa/core/tensor.h"
#include "singa/proto/io.pb.h"
#include "singa/proto/model.pb.h"

namespace singa {

class Transformer {
 public:
  Transformer() {}
  virtual ~Transformer() {}

  virtual void Setup(const TransformerConf& conf) {}

  virtual Tensor Apply(int flag, Tensor& input) = 0;
};

class ImageTransformer: public Transformer {
 public:
  void Setup(const TransformerConf& conf) override {
    featurewise_center_ = conf.featurewise_center();
    featurewise_std_norm_ = conf.featurewise_std_norm();
    resize_height_ = conf.resize_height();
    resize_width_ = conf.resize_width();
    rescale_ = conf.rescale();
    horizontal_mirror_ = conf.horizontal_mirror();
    image_dim_order_ = conf.image_dim_order();
    
    if (conf.crop_shape_size() == 2)
      crop_shape_ = {conf.crop_shape(0), conf.crop_shape(1)};      
  }

  Tensor Apply(int flag, Tensor& input) override;

  const bool featurewise_center() const { return featurewise_center_; }
  const bool featurewise_std_norm() const { return featurewise_std_norm_; }
  const bool horizontal_mirror() const { return horizontal_mirror_; }
  const int resize_height() const { return resize_height_; }
  const int resize_width() const { return resize_width_; }
  const float rescale() const { return rescale_; }
  const Shape crop_shape() const { return crop_shape_; }
  const string image_dim_order() const { return image_dim_order_; }

 private:
  bool featurewise_center_ = false;
  bool featurewise_std_norm_ = false;
  bool horizontal_mirror_ = false;
  int resize_height_ = 0;
  int resize_width_ = 0;
  float rescale_ = 0.f;
  Shape crop_shape_ = {};
  std::string image_dim_order_ = "CHW";
};

#ifdef USE_OPENCV
Tensor resize(Tensor& input, const size_t resize_height, 
         const size_t resize_width, const string& image_dim_order);
#endif
Tensor crop(Tensor& input, const size_t crop_height, 
             const size_t crop_width, const size_t crop_h_offset, 
             const size_t crop_w_offset, const string& image_dim_order);
Tensor mirror(Tensor& input, const bool horizontal_mirror, 
             const bool vertical_mirror, const string& image_dim_order);
} // namespace singa
#endif  // SINGA_IO_TRANSFORMER_H_