-
Notifications
You must be signed in to change notification settings - Fork 14
/
image.h
195 lines (176 loc) · 7.25 KB
/
image.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
// 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.
/** \file image.h
* \brief Optional image-specific helpers and specializations.
*/
#ifndef NDARRAY_IMAGE_H
#define NDARRAY_IMAGE_H
#include "array/array.h"
namespace nda {
/** A generic image is any 3D array with dimensions x, y, c. c represents the
* channels of the image, typically it will have extent 3 or 4, with red, green,
* and blue mapped to indices in this dimension. */
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>;
/** A 'chunky' image is an array with 3 dimensions x, y, c, where c is dense,
* and the dimension with the next stride is x. The stride in x may be larger
* than the number of channels, to allow for padding pixels to a convenient
* alignment. This is a common image storage format used by many programs
* working with images. */
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>;
/** Calls `fn` for each index in an image shape `s`. c is the innermost
* dimension of the loop nest. */
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);
}
}
};
/** A 'planar' image is an array with dimensions x, y, c, where x is dense. This
* format is less common, but more convenient for optimization, particularly
* SIMD vectorization. Note that this shape also supports 'line-chunky' storage
* orders. */
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 {
/** The result of the crop has min 0, 0. */
zero,
/** The result indices inside the crop are the same as the original indices,
* and the result indices outside the crop are out of bounds. */
crop,
};
/** Crop an image shape `s` to the indices [`x0`, `x1`) x [`y0`, `y1`). The origin of
* the new shape is determined by `origin`. */
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;
}
/** Crop the `im` image or image ref to the interval [`x0`, `x1`) x [`y0`, `y1`). The
* result is a ref of the input image. The origin of the result is determined by
* `origin`. */
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);
}
/** Get a 2-dimensional ref of the `Channel` channel of the `im` image or image
* ref. */
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