Camera tracking: merge hybrid tracker from tomato branch
[blender.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 #ifndef __APPLE__
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 #else // This is for Linux.
54   return memalign(alignment, size);
55 #endif
56 }
57
58 void aligned_free(void *ptr) {
59 #ifdef _WIN32
60   _aligned_free(ptr);
61 #else
62   free(ptr);
63 #endif
64 }
65
66 bool RegionIsInBounds(const FloatImage &image1,
67                       double x, double y,
68                       int half_window_size) {
69   // Check the minimum coordinates.
70   int min_x = floor(x) - half_window_size - 1;
71   int min_y = floor(y) - half_window_size - 1;
72   if (min_x < 0.0 ||
73       min_y < 0.0) {
74     return false;
75   }
76
77   // Check the maximum coordinates.
78   int max_x = ceil(x) + half_window_size + 1;
79   int max_y = ceil(y) + half_window_size + 1;
80   if (max_x > image1.cols() ||
81       max_y > image1.rows()) {
82     return false;
83   }
84
85   // Ok, we're good.
86   return true;
87 }
88
89 #ifdef __SSE2__
90
91 // Compute the sub of absolute differences between the arrays "a" and "b". 
92 // The array "a" is assumed to be 16-byte aligned, while "b" is not. The
93 // result is returned as the first and third elements of __m128i if
94 // interpreted as a 4-element 32-bit integer array. The SAD is the sum of the
95 // elements.
96 //
97 // The function requires size % 16 valid extra elements at the end of both "a"
98 // and "b", since the SSE load instructionst will pull in memory past the end
99 // of the arrays if their size is not a multiple of 16.
100 inline static __m128i SumOfAbsoluteDifferencesContiguousSSE(
101     const unsigned char *a,  // aligned
102     const unsigned char *b,  // not aligned
103     unsigned int size,
104     __m128i sad) {
105   // Do the bulk of the work as 16-way integer operations.
106   for(unsigned int j = 0; j < size / 16; j++) {
107     sad = _mm_add_epi32(sad, _mm_sad_epu8( _mm_load_si128 ((__m128i*)(a + 16 * j)),
108                                            _mm_loadu_si128((__m128i*)(b + 16 * j))));
109   }
110   // Handle the trailing end.
111   // TODO(keir): Benchmark to verify that the below SSE is a win compared to a
112   // hand-rolled loop. It's not clear that the hand rolled loop would be slower
113   // than the potential cache miss when loading the immediate table below.
114   //
115   // An alternative to this version is to take a packet of all 1's then do a
116   // 128-bit shift. The issue is that the shift instruction needs an immediate
117   // amount rather than a variable amount, so the branch instruction here must
118   // remain. See _mm_srli_si128 and  _mm_slli_si128.
119   unsigned int remainder = size % 16u;
120   if (remainder) {
121     unsigned int j = size / 16;
122     __m128i a_trail = _mm_load_si128 ((__m128i*)(a + 16 * j));
123     __m128i b_trail = _mm_loadu_si128((__m128i*)(b + 16 * j));
124     __m128i mask;
125     switch (remainder) {
126 #define X 0xff
127       case  1: mask = _mm_setr_epi8(X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
128       case  2: mask = _mm_setr_epi8(X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
129       case  3: mask = _mm_setr_epi8(X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
130       case  4: mask = _mm_setr_epi8(X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
131       case  5: mask = _mm_setr_epi8(X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
132       case  6: mask = _mm_setr_epi8(X, X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
133       case  7: mask = _mm_setr_epi8(X, X, X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0, 0); break;
134       case  8: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0, 0); break;
135       case  9: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, 0, 0, 0, 0, 0, 0, 0); break;
136       case 10: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, 0, 0, 0, 0, 0, 0); break;
137       case 11: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, 0, 0, 0, 0, 0); break;
138       case 12: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, X, 0, 0, 0, 0); break;
139       case 13: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, X, X, 0, 0, 0); break;
140       case 14: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, X, X, X, 0, 0); break;
141       case 15: mask = _mm_setr_epi8(X, X, X, X, X, X, X, X, X, X, X, X, X, X, X, 0); break;
142 #undef X
143     }
144     sad = _mm_add_epi32(sad, _mm_sad_epu8(_mm_and_si128(mask, a_trail),
145                                           _mm_and_si128(mask, b_trail)));
146   }
147   return sad;
148 }
149 #endif
150
151 // Computes the sum of absolute differences between pattern and image. Pattern
152 // must be 16-byte aligned, and the stride must be a multiple of 16. The image
153 // does pointer does not have to be aligned.
154 int SumOfAbsoluteDifferencesContiguousImage(
155     const unsigned char *pattern,
156     unsigned int pattern_width,
157     unsigned int pattern_height,
158     unsigned int pattern_stride,
159     const unsigned char *image,
160     unsigned int image_stride) {
161 #ifdef __SSE2__
162   // TODO(keir): Add interleaved accumulation, where accumulation is done into
163   // two or more SSE registers that then get combined at the end. This reduces
164   // instruction dependency; in Eigen's squared norm code, splitting the
165   // accumulation produces a ~2x speedup. It's not clear it will help here,
166   // where the number of SSE instructions in the inner loop is smaller.
167   __m128i sad = _mm_setzero_si128();
168   for (int r = 0; r < pattern_height; ++r) {
169     sad = SumOfAbsoluteDifferencesContiguousSSE(&pattern[pattern_stride * r],
170                                                 &image[image_stride * r],
171                                                 pattern_width,
172                                                 sad);
173   }
174   return _mm_cvtsi128_si32(
175              _mm_add_epi32(sad,
176                  _mm_shuffle_epi32(sad, _MM_SHUFFLE(3, 0, 1, 2))));
177 #else
178   int sad = 0;
179   for (int r = 0; r < pattern_height; ++r) {
180     for (int c = 0; c < pattern_width; ++c) {
181       sad += abs(pattern[pattern_stride * r + c] - image[image_stride * r + c]);
182     }
183   }
184   return sad;
185 #endif
186 }
187
188 // Sample a region of size width, height centered at x,y in image, converting
189 // from float to byte in the process. Samples from the first channel. Puts
190 // result into *pattern.
191 void SampleRectangularPattern(const FloatImage &image,
192                               double x, double y,
193                               int width,
194                               int height,
195                               int pattern_stride,
196                               unsigned char *pattern) {
197   // There are two cases for width and height: even or odd. If it's odd, then
198   // the bounds [-width / 2, width / 2] works as expected. However, for even,
199   // this results in one extra access past the end. So use < instead of <= in
200   // the loops below, but increase the end limit by one in the odd case.
201   int end_width = (width / 2) + (width % 2);
202   int end_height = (height / 2) + (height % 2);
203   for (int r = -height / 2; r < end_height; ++r) {
204     for (int c = -width / 2; c < end_width; ++c) {
205       pattern[pattern_stride * (r + height / 2) +  c + width / 2] =
206           SampleLinear(image, y + r, x + c, 0) * 255.0;
207     }
208   }
209 }
210
211 // Returns x rounded up to the nearest multiple of alignment.
212 inline int PadToAlignment(int x, int alignment) {
213   if (x % alignment != 0) {
214     x += alignment - (x % alignment);
215   }
216   return x;
217 }
218
219 // Sample a region centered at x,y in image with size extending by half_width
220 // from x. Samples from the first channel. The resulting array is placed in
221 // *pattern, and the stride, which will be a multiple of 16 if SSE is enabled,
222 // is returned in *pattern_stride.
223 //
224 // NOTE: Caller must free *pattern with aligned_malloc() from above.
225 void SampleSquarePattern(const FloatImage &image,
226                          double x, double y,
227                          int half_width,
228                          unsigned char **pattern,
229                          int *pattern_stride) {
230   int width = 2 * half_width + 1;
231   // Allocate an aligned block with padding on the end so each row of the
232   // pattern starts on a 16-byte boundary.
233   *pattern_stride = PadToAlignment(width, 16);
234   int pattern_size_bytes = *pattern_stride * width;
235   *pattern = static_cast<unsigned char *>(
236       aligned_malloc(pattern_size_bytes, 16));
237   SampleRectangularPattern(image, x, y, width, width,
238                            *pattern_stride,
239                            *pattern);
240 }
241
242 // NOTE: Caller must free *image with aligned_malloc() from above.
243 void FloatArrayToByteArrayWithPadding(const FloatImage &float_image,
244                                       unsigned char **image,
245                                       int *image_stride) {
246   // Allocate enough so that accessing 16 elements past the end is fine.
247   *image_stride = float_image.Width() + 16;
248   *image = static_cast<unsigned char *>(
249       aligned_malloc(*image_stride * float_image.Height(), 16));
250   for (int i = 0; i < float_image.Height(); ++i) {
251     for (int j = 0; j < float_image.Width(); ++j) {
252       (*image)[*image_stride * i + j] =
253           static_cast<unsigned char>(255.0 * float_image(i, j, 0));
254     }
255   }
256 }
257
258 }  // namespace
259
260 // TODO(keir): Compare the "sharpness" of the peak around the best pixel. It's
261 // probably worth plotting a few examples to see what the histogram of SAD
262 // values for every hypothesis looks like.
263 //
264 // TODO(keir): Priority queue for multiple hypothesis.
265 bool BruteRegionTracker::Track(const FloatImage &image1,
266                                const FloatImage &image2,
267                                double  x1, double  y1,
268                                double *x2, double *y2) const {
269   if (!RegionIsInBounds(image1, x1, y1, half_window_size)) {
270     LG << "Fell out of image1's window with x1=" << x1 << ", y1=" << y1
271        << ", hw=" << half_window_size << ".";
272     return false;
273   }
274   int pattern_width = 2 * half_window_size + 1;
275
276   Array3Df image_and_gradient1;
277   Array3Df image_and_gradient2;
278   BlurredImageAndDerivativesChannels(image1, 0.9, &image_and_gradient1);
279   BlurredImageAndDerivativesChannels(image2, 0.9, &image_and_gradient2);
280
281   // Sample the pattern to get it aligned to an image grid.
282   unsigned char *pattern;
283   int pattern_stride;
284   SampleSquarePattern(image_and_gradient1, x1, y1, half_window_size,
285                       &pattern,
286                       &pattern_stride);
287
288   // Convert the search area directly to bytes without sampling.
289   unsigned char *search_area;
290   int search_area_stride;
291   FloatArrayToByteArrayWithPadding(image_and_gradient2, &search_area, &search_area_stride);
292
293   // Try all possible locations inside the search area. Yes, everywhere.
294   int best_i, best_j, best_sad = INT_MAX;
295   for (int i = 0; i < image2.Height() - pattern_width; ++i) {
296     for (int j = 0; j < image2.Width() - pattern_width; ++j) {
297       int sad = SumOfAbsoluteDifferencesContiguousImage(pattern,
298                                                         pattern_width,
299                                                         pattern_width,
300                                                         pattern_stride,
301                                                         search_area + search_area_stride * i + j,
302                                                         search_area_stride);
303       if (sad < best_sad) {
304         best_i = i;
305         best_j = j;
306         best_sad = sad;
307       }
308     }
309   }
310
311   aligned_free(pattern);
312   aligned_free(search_area);
313
314   if (best_sad != INT_MAX) {
315     *x2 = best_j + half_window_size;
316     *y2 = best_i + half_window_size;
317     return true;
318   }
319   return false;
320 }
321
322 }  // namespace libmv