image.h

Go to the documentation of this file.

// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef NDARRAY_IMAGE_H
#define NDARRAY_IMAGE_H

#include "array/array.h"

namespace nda {

using image_shape = shape<dim<>, dim<>, dim<>>;
template <class T, class Alloc = std::allocator<T>>
using image = array<T, image_shape, Alloc>;
template <class T>
using image_ref = array_ref<T, image_shape>;
template <class T>
using const_image_ref = image_ref<const T>;

template <index_t Channels = dynamic, index_t XStride = Channels>
using chunky_image_shape = shape<strided_dim<XStride>, dim<>, dense_dim<0, Channels>>;
template <class T, index_t Channels = dynamic, index_t XStride = Channels,
    class Alloc = std::allocator<T>>
using chunky_image = array<T, chunky_image_shape<Channels, XStride>, Alloc>;
template <class T, index_t Channels = dynamic, index_t XStride = Channels>
using chunky_image_ref = array_ref<T, chunky_image_shape<Channels, XStride>>;
template <class T, index_t Channels = dynamic, index_t XStride = Channels>
using const_chunky_image_ref = chunky_image_ref<const T, Channels, XStride>;

template <class Shape, class Fn>
void for_each_image_index(const Shape& s, Fn&& fn) {
  // Images should always be iterated with c as the innermost loop. Even when
  // the image is planar, the number of channels is generally small, and many
  // operations use all of the channels at the same time (all-to-all
  // communication in the c dimension).
  for (index_t y : s.y()) {
    for (index_t x : s.x()) {
      for (index_t c : s.c()) {
        fn(std::tuple<index_t, index_t, index_t>(x, y, c));
      }
    }
  }
}

template <index_t Channels, index_t XStride>
class shape_traits<chunky_image_shape<Channels, XStride>> {
public:
  typedef chunky_image_shape<Channels, XStride> shape_type;

  template <class Fn>
  static void for_each_index(const shape_type& s, Fn&& fn) {
    for_each_image_index(s, fn);
  }

  template <class Ptr, class Fn>
  static void for_each_value(const shape_type& s, Ptr base, Fn&& fn) {
    for_each_image_index(
        s, [=, fn = std::move(fn)](const typename shape_type::index_type& i) { fn(base[s[i]]); });
  }
};

template <index_t Channels>
class shape_traits<chunky_image_shape<Channels>> {
public:
  typedef chunky_image_shape<Channels> shape_type;

  template <class Fn>
  static void for_each_index(const shape_type& s, Fn&& fn) {
    for_each_image_index(s, fn);
  }

  // When Channels == XStride, we can implement for_each_value by fusing
  // the x and c dimensions.
  template <class Ptr, class Fn>
  static void for_each_value(const shape_type& s, Ptr base, Fn&& fn) {
    const index_t channels = s.c().extent();
    if (channels == s.x().stride()) {
      dense_shape<2> opt_s({s.x().min() * channels, s.x().extent() * channels}, s.y());
      for_each_value_in_order(opt_s, base, fn);
    } else {
      for_each_value_in_order(s, base, fn);
    }
  }
};

using planar_image_shape = shape<dense_dim<>, dim<>, dim<>>;
template <class T, class Alloc = std::allocator<T>>
using planar_image = array<T, planar_image_shape, Alloc>;
template <class T>
using planar_image_ref = array_ref<T, planar_image_shape>;
template <class T>
using const_planar_image_ref = planar_image_ref<const T>;

enum class crop_origin {
  zero,
  crop,
};

template <class Shape>
Shape crop_image_shape(Shape s, index_t x0, index_t y0, index_t x1, index_t y1,
    crop_origin origin = crop_origin::crop) {
  s.x().set_extent(x1 - x0);
  s.y().set_extent(y1 - y0);
  switch (origin) {
  case crop_origin::zero:
    s.x().set_min(0);
    s.y().set_min(0);
    break;
  case crop_origin::crop:
    s.x().set_min(x0);
    s.y().set_min(y0);
    break;
  }
  return s;
}

template <class T, class Shape>
array_ref<T, Shape> crop(const array_ref<T, Shape>& im, index_t x0, index_t y0, index_t x1,
    index_t y1, crop_origin origin = crop_origin::crop) {
  x0 = internal::max(x0, im.x().min());
  y0 = internal::max(y0, im.y().min());
  x1 = internal::min(x1, im.x().max() + 1);
  y1 = internal::min(y1, im.y().max() + 1);
  Shape cropped_shape = crop_image_shape(im.shape(), x0, y0, x1, y1, origin);
  index_t c0 = im.c().min();
  T* base = im.base() != nullptr ? &im(x0, y0, c0) : nullptr;
  if (origin == crop_origin::crop) {
    base = internal::pointer_add(base, -cropped_shape(x0, y0, c0));
  }
  return array_ref<T, Shape>(base, cropped_shape);
}
template <class T, class Shape>
array_ref<const T, Shape> crop(const array<T, Shape>& im, index_t x0, index_t y0, index_t x1,
    index_t y1, crop_origin origin = crop_origin::crop) {
  return crop(im.ref(), x0, y0, x1, y1, origin);
}
template <class T, class Shape>
array_ref<T, Shape> crop(array<T, Shape>& im, index_t x0, index_t y0, index_t x1, index_t y1,
    crop_origin origin = crop_origin::crop) {
  return crop(im.ref(), x0, y0, x1, y1, origin);
}

template <class T, class Shape>
auto slice_channel(const array_ref<T, Shape>& im, index_t channel) {
  return im(im.x(), im.y(), channel);
}
template <class T, class Shape, class Alloc>
auto slice_channel(const array<T, Shape, Alloc>& im, index_t channel) {
  return slice_channel(im.ref(), channel);
}
template <class T, class Shape, class Alloc>
auto slice_channel(array<T, Shape, Alloc>& im, index_t channel) {
  return slice_channel(im.ref(), channel);
}

} // namespace nda

#endif // NDARRAY_IMAGE_H