Tensor mirror()

in src/io/image_transformer.cc [273:352]

• 78 lines of code
• 24 McCabe index (conditional complexity)

Tensor mirror(Tensor& input, const bool horizontal_mirror,
              const bool vertical_mirror, const string& image_dim_order) {
  CHECK_LE(input.nDim(), 4u);
  CHECK_GE(input.nDim(), 2u);
  if (!horizontal_mirror && !vertical_mirror) return input;

  Tensor output;
  const float* in = input.data<float>();
  size_t out_idx = 0, in_idx = 0;
  if (input.nDim() == 4u) {
    /// TODO
    LOG(FATAL) << "Not implemented";
  } else if (input.nDim() == 3u) {
    if (image_dim_order == "CHW") {
      size_t height = input.shape(1), width = input.shape(2),
             channel = input.shape(0);
      float* out = new float[height * width * channel];
      for (size_t c = 0; c < channel; c++) {
        for (size_t h = 0; h < height; h++) {
          for (size_t w = 0; w < width; w++) {
            in_idx = (c * height + h) * width + w;
            if (horizontal_mirror && vertical_mirror)
              out_idx = (c * height + (height - 1 - h)) * width + (width - 1 - w);
            else if (horizontal_mirror)
              out_idx = (c * height + h) * width + (width - 1 - w);
            else /// only do vertical mirror
              out_idx = (c * height + (height - 1 - h)) * width + w;
            out[out_idx] = in[in_idx];
          }
        }
      }
      output.Resize(Shape{channel, height, width});
      output.CopyDataFromHostPtr<float>(out, height * width * channel);
      delete[] out;
    } else if (image_dim_order == "HWC") {
      size_t height = input.shape(0), width = input.shape(1),
             channel = input.shape(2);
      float* out = new float[height * width * channel];
      for (size_t c = 0; c < channel; c++) {
        for (size_t h = 0; h < height; h++) {
          for (size_t w = 0; w < width; w++) {
            in_idx = (h * width + w) * channel + c;
            if (horizontal_mirror && vertical_mirror)
              out_idx = ((height - 1 - h) * width + (width - 1 - w)) * channel + c;
            else if (horizontal_mirror)
              out_idx = (h * width + (width - 1 - w)) * channel + c;
            else /// only do vertical mirror
              out_idx = ((height - 1 - h) * width + w) * channel + c;
            out[out_idx] = in[in_idx];
          }
        }
      }
      output.Resize(Shape{height, width, channel});
      output.CopyDataFromHostPtr<float>(out, height * width * channel);
      delete[] out;
    } else {
      LOG(FATAL) << "Unknow dimension order for images " << image_dim_order
                 << " Only support 'HWC' and 'CHW'";
    }
  } else { /// 2D gray image
    size_t height = input.shape(0), width = input.shape(1);
    float* out = new float[height * width];
    for (size_t h = 0; h < height; h++) {
      for (size_t w = 0; w < width; w++) {
        in_idx = h * width + w;
        if (horizontal_mirror && vertical_mirror)
          out_idx = (height - 1 - h) * width + (width - 1 - w);
        else if (horizontal_mirror)
          out_idx = h * width + (width - 1 - w);
        else /// only do vertical mirror
          out_idx = (height - 1 - h) * width + w;
        out[out_idx] = in[in_idx];
      }
    }
    output.Resize(Shape{height, width});
    output.CopyDataFromHostPtr<float>(out, height * width);
    delete[] out;
  }
  return output;
}