Merging r42945 through r43024 from trunk into soc-2011-tomato
[blender-staging.git] / extern / libmv / libmv / tracking / brute_region_tracker.cc
1 // Copyright (c) 2011 libmv authors.
2 //
3 // Permission is hereby granted, free of charge, to any person obtaining a copy
4 // of this software and associated documentation files (the "Software"), to
5 // deal in the Software without restriction, including without limitation the
6 // rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
7 // sell copies of the Software, and to permit persons to whom the Software is
8 // furnished to do so, subject to the following conditions:
9 //
10 // The above copyright notice and this permission notice shall be included in
11 // all copies or substantial portions of the Software.
12 //
13 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
19 // IN THE SOFTWARE.
20
21 #include "libmv/tracking/brute_region_tracker.h"
22
23 #ifdef __SSE2__
24 #include <emmintrin.h>
25 #endif
26
27 #if !defined(__APPLE__) && !defined(__FreeBSD__)
28 // Needed for memalign on Linux and _aligned_alloc on Windows.
29 #include <malloc.h>
30 #else
31 // Apple's malloc is 16-byte aligned, and does not have malloc.h, so include
32 // stdilb instead.
33 #include <cstdlib>
34 #endif
35
36 #include "libmv/image/image.h"
37 #include "libmv/image/convolve.h"
38 #include "libmv/image/sample.h"
39 #include "libmv/logging/logging.h"
40
41 namespace libmv {
42 namespace {
43
44 // TODO(keir): It's stupid that this is needed here. Push this somewhere else.
45 void *aligned_malloc(int size, int alignment) {
46 #ifdef _WIN32
47   return _aligned_malloc(size, alignment);
48 #elif __APPLE__
49   // On Mac OS X, both the heap and the stack are guaranteed 16-byte aligned so
50   // they work natively with SSE types with no further work.
51   CHECK_EQ(alignment, 16);
52   return malloc(size);
53 #elif __FreeBSD__
54   void *result;
55
56   if(posix_memalign(&result, alignment, size)) {
57     // non-zero means allocation error
58     // either no allocation or bad alignment value
59     return NULL;
60   }
61   return result;
62 #else // This is for Linux.
63   return memalign(alignment, size);
64 #endif
65 }
66
67 void aligned_free(void *ptr) {
68 #ifdef _WIN32
69   _aligned_free(ptr);
70 #else
71   free(ptr);
72 #endif
73 }
74
75 bool RegionIsInBounds(const FloatImage &image1,
76                       double x, double y,
77                       int half_window_size) {
78   // Check the minimum coordinates.
79   int min_x = floor(x) - half_window_size - 1;
80   int min_y = floor(y) - half_window_size - 1;
81   if (min_x < 0.0 ||
82       min_y < 0.0) {
83     return false;
84   }
85
86   // Check the maximum coordinates.
87   int max_x = ceil(x) + half_window_size + 1;
88   int max_y = ceil(y) + half_window_size + 1;
89   if (max_x > image1.cols() ||
90       max_y > image1.rows()) {
91     return false;
92   }
93
94   // Ok, we're good.
95   return true;
96 }
97
98 #ifdef __SSE2__
99
100 // Compute the sub of absolute differences between the arrays "a" and "b". 
101 // The array "a" is assumed to be 16-byte aligned, while "b" is not. The
102 // result is returned as the first and third elements of __m128i if
103 // interpreted as a 4-element 32-bit integer array. The SAD is the sum of the
104 // elements.
105 //
106 // The function requires size % 16 valid extra elements at the end of both "a"
107 // and "b", since the SSE load instructionst will pull in memory past the end
108 // of the arrays if their size is not a multiple of 16.
109 inline static __m128i SumOfAbsoluteDifferencesContiguousSSE(
110     const unsigned char *a,  // aligned
111     const unsigned char *b,  // not aligned
112     unsigned int size,
113     __m128i sad) {
114   // Do the bulk of the work as 16-way integer operations.
115   for(unsigned int j = 0; j < size / 16; j++) {
116     sad = _mm_add_epi32(sad, _mm_sad_epu8( _mm_load_si128 ((__m128i*)(a + 16 * j)),
117                                            _mm_loadu_si128((__m128i*)(b + 16 * j))));
118   }
119   // Handle the trailing end.
120   // TODO(keir): Benchmark to verify that the below SSE is a win compared to a
121   // hand-rolled loop. It's not clear that the hand rolled loop would be slower
122   // than the potential cache miss when loading the immediate table below.
123   //
124   // An alternative to this version is to take a packet of all 1's then do a
125   // 128-bit shift. The issue is that the shift instruction needs an immediate
126   // amount rather than a variable amount, so the branch instruction here must
127   // remain. See _mm_srli_si128 and  _mm_slli_si128.
128   unsigned int remainder = size % 16u;
129   if (remainder) {
130     unsigned int j = size / 16;
131     __m128i a_trail = _mm_load_si128 ((__m128i*)(a + 16 * j));
132     __m128i b_trail = _mm_loadu_si128((__m128i*)(b + 16 * j));
133     __m128i mask;
134     switch (remainder) {
135 #define X 0xff
136       case  1: mask = _mm_setr_epi8(X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
137       case  2: mask = _mm_setr_epi8(X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
138       case  3: mask = _mm_setr_epi8(X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
139       case  4: mask = _mm_setr_epi8(X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
140       case  5: mask = _mm_setr_epi8(X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
141       case  6: mask = _mm_setr_epi8(X, X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
142       case  7: mask = _mm_setr_epi8(X, X, X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
143       case  8: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0); break;
144       case  9: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0); break;
145       case 10: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, 0, 0, 0, 0, 0, 0); break;
146       case 11: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, 0, 0, 0, 0, 0); break;
147       case 12: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, X, 0, 0, 0, 0); break;
148       case 13: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, X, X, 0, 0, 0); break;
149       case 14: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, X, X, X, 0, 0); break;
150       case 15: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, X, X, X, X, 0); break;
151 #undef X
152     }
153     sad = _mm_add_epi32(sad, _mm_sad_epu8(_mm_and_si128(mask, a_trail),
154                                           _mm_and_si128(mask, b_trail)));
155   }
156   return sad;
157 }
158 #endif
159
160 // Computes the sum of absolute differences between pattern and image. Pattern
161 // must be 16-byte aligned, and the stride must be a multiple of 16. The image
162 // does pointer does not have to be aligned.
163 int SumOfAbsoluteDifferencesContiguousImage(
164     const unsigned char *pattern,
165     unsigned int pattern_width,
166     unsigned int pattern_height,
167     unsigned int pattern_stride,
168     const unsigned char *image,
169     unsigned int image_stride) {
170 #ifdef __SSE2__
171   // TODO(keir): Add interleaved accumulation, where accumulation is done into
172   // two or more SSE registers that then get combined at the end. This reduces
173   // instruction dependency; in Eigen's squared norm code, splitting the
174   // accumulation produces a ~2x speedup. It's not clear it will help here,
175   // where the number of SSE instructions in the inner loop is smaller.
176   __m128i sad = _mm_setzero_si128();
177   for (int r = 0; r < pattern_height; ++r) {
178     sad = SumOfAbsoluteDifferencesContiguousSSE(&pattern[pattern_stride * r],
179                                                 &image[image_stride * r],
180                                                 pattern_width,
181                                                 sad);
182   }
183   return _mm_cvtsi128_si32(
184              _mm_add_epi32(sad,
185                  _mm_shuffle_epi32(sad, _MM_SHUFFLE(3, 0, 1, 2))));
186 #else
187   int sad = 0;
188   for (int r = 0; r < pattern_height; ++r) {
189     for (int c = 0; c < pattern_width; ++c) {
190       sad += abs(pattern[pattern_stride * r + c] - image[image_stride * r + c]);
191     }
192   }
193   return sad;
194 #endif
195 }
196
197 // Sample a region of size width, height centered at x,y in image, converting
198 // from float to byte in the process. Samples from the first channel. Puts
199 // result into *pattern.
200 void SampleRectangularPattern(const FloatImage &image,
201                               double x, double y,
202                               int width,
203                               int height,
204                               int pattern_stride,
205                               unsigned char *pattern) {
206   // There are two cases for width and height: even or odd. If it's odd, then
207   // the bounds [-width / 2, width / 2] works as expected. However, for even,
208   // this results in one extra access past the end. So use < instead of <= in
209   // the loops below, but increase the end limit by one in the odd case.
210   int end_width = (width / 2) + (width % 2);
211   int end_height = (height / 2) + (height % 2);
212   for (int r = -height / 2; r < end_height; ++r) {
213     for (int c = -width / 2; c < end_width; ++c) {
214       pattern[pattern_stride * (r + height / 2) +  c + width / 2] =
215           SampleLinear(image, y + r, x + c, 0) * 255.0;
216     }
217   }
218 }
219
220 // Returns x rounded up to the nearest multiple of alignment.
221 inline int PadToAlignment(int x, int alignment) {
222   if (x % alignment != 0) {
223     x += alignment - (x % alignment);
224   }
225   return x;
226 }
227
228 // Sample a region centered at x,y in image with size extending by half_width
229 // from x. Samples from the first channel. The resulting array is placed in
230 // *pattern, and the stride, which will be a multiple of 16 if SSE is enabled,
231 // is returned in *pattern_stride.
232 //
233 // NOTE: Caller must free *pattern with aligned_malloc() from above.
234 void SampleSquarePattern(const FloatImage &image,
235                          double x, double y,
236                          int half_width,
237                          unsigned char **pattern,
238                          int *pattern_stride) {
239   int width = 2 * half_width + 1;
240   // Allocate an aligned block with padding on the end so each row of the
241   // pattern starts on a 16-byte boundary.
242   *pattern_stride = PadToAlignment(width, 16);
243   int pattern_size_bytes = *pattern_stride * width;
244   *pattern = static_cast<unsigned char *>(
245       aligned_malloc(pattern_size_bytes, 16));
246   SampleRectangularPattern(image, x, y, width, width,
247                            *pattern_stride,
248                            *pattern);
249 }
250
251 // NOTE: Caller must free *image with aligned_malloc() from above.
252 void FloatArrayToByteArrayWithPadding(const FloatImage &float_image,
253                                       unsigned char **image,
254                                       int *image_stride) {
255   // Allocate enough so that accessing 16 elements past the end is fine.
256   *image_stride = float_image.Width() + 16;
257   *image = static_cast<unsigned char *>(
258       aligned_malloc(*image_stride * float_image.Height(), 16));
259   for (int i = 0; i < float_image.Height(); ++i) {
260     for (int j = 0; j < float_image.Width(); ++j) {
261       (*image)[*image_stride * i + j] =
262           static_cast<unsigned char>(255.0 * float_image(i, j, 0));
263     }
264   }
265 }
266
267 }  // namespace
268
269 // TODO(keir): Compare the "sharpness" of the peak around the best pixel. It's
270 // probably worth plotting a few examples to see what the histogram of SAD
271 // values for every hypothesis looks like.
272 //
273 // TODO(keir): Priority queue for multiple hypothesis.
274 bool BruteRegionTracker::Track(const FloatImage &image1,
275                                const FloatImage &image2,
276                                double  x1, double  y1,
277                                double *x2, double *y2) const {
278   if (!RegionIsInBounds(image1, x1, y1, half_window_size)) {
279     LG << "Fell out of image1's window with x1=" << x1 << ", y1=" << y1
280        << ", hw=" << half_window_size << ".";
281     return false;
282   }
283   int pattern_width = 2 * half_window_size + 1;
284
285   Array3Df image_and_gradient1;
286   Array3Df image_and_gradient2;
287   BlurredImageAndDerivativesChannels(image1, 0.9, &image_and_gradient1);
288   BlurredImageAndDerivativesChannels(image2, 0.9, &image_and_gradient2);
289
290   // Sample the pattern to get it aligned to an image grid.
291   unsigned char *pattern;
292   int pattern_stride;
293   SampleSquarePattern(image_and_gradient1, x1, y1, half_window_size,
294                       &pattern,
295                       &pattern_stride);
296
297   // Convert the search area directly to bytes without sampling.
298   unsigned char *search_area;
299   int search_area_stride;
300   FloatArrayToByteArrayWithPadding(image_and_gradient2, &search_area, &search_area_stride);
301
302   // Try all possible locations inside the search area. Yes, everywhere.
303   int best_i, best_j, best_sad = INT_MAX;
304   for (int i = 0; i < image2.Height() - pattern_width; ++i) {
305     for (int j = 0; j < image2.Width() - pattern_width; ++j) {
306       int sad = SumOfAbsoluteDifferencesContiguousImage(pattern,
307                                                         pattern_width,
308                                                         pattern_width,
309                                                         pattern_stride,
310                                                         search_area + search_area_stride * i + j,
311                                                         search_area_stride);
312       if (sad < best_sad) {
313         best_i = i;
314         best_j = j;
315         best_sad = sad;
316       }
317     }
318   }
319
320   aligned_free(pattern);
321   aligned_free(search_area);
322
323   if (best_sad != INT_MAX) {
324     *x2 = best_j + half_window_size;
325     *y2 = best_i + half_window_size;
326     return true;
327   }
328   return false;
329 }
330
331 }  // namespace libmv