4a02031bfe6f8b6e660cbf9cfefad943a193dc09
[blender.git] / extern / libmv / libmv / autotrack / autotrack.cc
1 // Copyright (c) 2014 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 // Author: mierle@gmail.com (Keir Mierle)
22
23 #include "libmv/autotrack/autotrack.h"
24 #include "libmv/autotrack/quad.h"
25 #include "libmv/autotrack/frame_accessor.h"
26 #include "libmv/autotrack/predict_tracks.h"
27 #include "libmv/base/scoped_ptr.h"
28 #include "libmv/logging/logging.h"
29 #include "libmv/numeric/numeric.h"
30
31 namespace mv {
32
33 namespace {
34
35 class DisableChannelsTransform : public FrameAccessor::Transform {
36  public:
37   DisableChannelsTransform(int visible_channels)
38       : visible_channels_(visible_channels) {  }
39
40   int64_t key() const {
41     // We invert bits here so key is never null.
42     return !visible_channels_;
43   }
44
45   void run(const FloatImage& input, FloatImage* output) const {
46     bool disable_red   = (visible_channels_ & Marker::CHANNEL_R) == 0,
47          disable_green = (visible_channels_ & Marker::CHANNEL_G) == 0,
48          disable_blue  = (visible_channels_ & Marker::CHANNEL_B) == 0;
49
50     LG << "Disabling channels: "
51        << (disable_red   ? "R " : "")
52        << (disable_green ? "G " : "")
53        << (disable_blue  ? "B" : "");
54
55     // It's important to rescale the resultappropriately so that e.g. if only
56     // blue is selected, it's not zeroed out.
57     float scale = (disable_red   ? 0.0f : 0.2126f) +
58                   (disable_green ? 0.0f : 0.7152f) +
59                   (disable_blue  ? 0.0f : 0.0722f);
60
61     output->Resize(input.Width(), input.Height(), 1);
62     for (int y = 0; y < input.Height(); y++) {
63       for (int x = 0; x < input.Width(); x++) {
64         float r = disable_red   ? 0.0f : input(x, y, 0);
65         float g = disable_green ? 0.0f : input(x, y, 1);
66         float b = disable_blue  ? 0.0f : input(x, y, 2);
67         (*output)(x, y, 0) = (0.2126f * r + 0.7152f * g + 0.0722f * b) / scale;
68       }
69     }
70   }
71
72  private:
73   // Bitfield representing visible channels, bits are from Marker::Channel.
74   int visible_channels_;
75 };
76
77 template<typename QuadT, typename ArrayT>
78 void QuadToArrays(const QuadT& quad, ArrayT* x, ArrayT* y) {
79   for (int i = 0; i < 4; ++i) {
80     x[i] = quad.coordinates(i, 0);
81     y[i] = quad.coordinates(i, 1);
82   }
83 }
84
85 void MarkerToArrays(const Marker& marker, double* x, double* y) {
86   Quad2Df offset_quad = marker.patch;
87   Vec2f origin = marker.search_region.Rounded().min;
88   offset_quad.coordinates.rowwise() -= origin.transpose();
89   QuadToArrays(offset_quad, x, y);
90   x[4] = marker.center.x() - origin(0);
91   y[4] = marker.center.y() - origin(1);
92 }
93
94 FrameAccessor::Key GetImageForMarker(const Marker& marker,
95                                      FrameAccessor* frame_accessor,
96                                      FloatImage* image) {
97   // TODO(sergey): Currently we pass float region to the accessor,
98   // but we don't want the accessor to decide the rounding, so we
99   // do rounding here.
100   // Ideally we would need to pass IntRegion to the frame accessor.
101   Region region = marker.search_region.Rounded();
102   libmv::scoped_ptr<FrameAccessor::Transform> transform = NULL;
103   if (marker.visible_channels != (Marker::CHANNEL_R |
104                                   Marker::CHANNEL_G |
105                                   Marker::CHANNEL_B)) {
106     transform.reset(new DisableChannelsTransform(marker.visible_channels));
107   }
108   return frame_accessor->GetImage(marker.clip,
109                                   marker.frame,
110                                   FrameAccessor::MONO,
111                                   0,  // No downscale for now.
112                                   &region,
113                                   transform.get(),
114                                   image);
115 }
116
117 }  // namespace
118
119 bool AutoTrack::TrackMarker(Marker* tracked_marker,
120                             TrackRegionResult* result,
121                             const TrackRegionOptions* track_options) {
122   // Try to predict the location of the second marker.
123   bool predicted_position = false;
124   if (PredictMarkerPosition(tracks_, tracked_marker)) {
125     LG << "Succesfully predicted!";
126     predicted_position = true;
127   } else {
128     LG << "Prediction failed; trying to track anyway.";
129   }
130
131   Marker reference_marker;
132   tracks_.GetMarker(tracked_marker->reference_clip,
133                     tracked_marker->reference_frame,
134                     tracked_marker->track,
135                     &reference_marker);
136
137   // Convert markers into the format expected by TrackRegion.
138   double x1[5], y1[5];
139   MarkerToArrays(reference_marker, x1, y1);
140
141   double x2[5], y2[5];
142   MarkerToArrays(*tracked_marker, x2, y2);
143
144   // TODO(keir): Technically this could take a smaller slice from the source
145   // image instead of taking one the size of the search window.
146   FloatImage reference_image;
147   FrameAccessor::Key reference_key = GetImageForMarker(reference_marker,
148                                                        frame_accessor_,
149                                                        &reference_image);
150   if (!reference_key) {
151     LG << "Couldn't get frame for reference marker: " << reference_marker;
152     return false;
153   }
154
155   FloatImage tracked_image;
156   FrameAccessor::Key tracked_key = GetImageForMarker(*tracked_marker,
157                                                      frame_accessor_,
158                                                      &tracked_image);
159   if (!tracked_key) {
160     LG << "Couldn't get frame for tracked marker: " << tracked_marker;
161     return false;
162   }
163
164   // Store original position befoer tracking, so we can claculate offset later.
165   Vec2f original_center = tracked_marker->center;
166
167   // Do the tracking!
168   TrackRegionOptions local_track_region_options;
169   if (track_options) {
170     local_track_region_options = *track_options;
171   }
172   local_track_region_options.num_extra_points = 1;  // For center point.
173   local_track_region_options.attempt_refine_before_brute = predicted_position;
174   TrackRegion(reference_image,
175               tracked_image,
176               x1, y1,
177               local_track_region_options,
178               x2, y2,
179               result);
180
181   // Copy results over the tracked marker.
182   Vec2f tracked_origin = tracked_marker->search_region.Rounded().min;
183   for (int i = 0; i < 4; ++i) {
184     tracked_marker->patch.coordinates(i, 0) = x2[i] + tracked_origin[0];
185     tracked_marker->patch.coordinates(i, 1) = y2[i] + tracked_origin[1];
186   }
187   tracked_marker->center(0) = x2[4] + tracked_origin[0];
188   tracked_marker->center(1) = y2[4] + tracked_origin[1];
189   Vec2f delta = tracked_marker->center - original_center;
190   tracked_marker->search_region.Offset(delta);
191   tracked_marker->source = Marker::TRACKED;
192   tracked_marker->status = Marker::UNKNOWN;
193   tracked_marker->reference_clip  = reference_marker.clip;
194   tracked_marker->reference_frame = reference_marker.frame;
195
196   // Release the images from the accessor cache.
197   frame_accessor_->ReleaseImage(reference_key);
198   frame_accessor_->ReleaseImage(tracked_key);
199
200   // TODO(keir): Possibly the return here should get removed since the results
201   // are part of TrackResult. However, eventually the autotrack stuff will have
202   // extra status (e.g. prediction fail, etc) that should get included.
203   return true;
204 }
205
206 void AutoTrack::AddMarker(const Marker& marker) {
207   tracks_.AddMarker(marker);
208 }
209
210 void AutoTrack::SetMarkers(vector<Marker>* markers) {
211   tracks_.SetMarkers(markers);
212 }
213
214 bool AutoTrack::GetMarker(int clip, int frame, int track,
215                           Marker* markers) const {
216   return tracks_.GetMarker(clip, frame, track, markers);
217 }
218
219 void AutoTrack::DetectAndTrack(const DetectAndTrackOptions& options) {
220   int num_clips = frame_accessor_->NumClips();
221   for (int clip = 0; clip < num_clips; ++clip) {
222     int num_frames = frame_accessor_->NumFrames(clip);
223     vector<Marker> previous_frame_markers;
224     // Q: How to decide track #s when detecting?
225     // Q: How to match markers from previous frame? set of prev frame tracks?
226     // Q: How to decide what markers should get tracked and which ones should not?
227     for (int frame = 0; frame < num_frames; ++frame) {
228       if (Cancelled()) {
229         LG << "Got cancel message while detecting and tracking...";
230         return;
231       }
232       // First, get or detect markers for this frame.
233       vector<Marker> this_frame_markers;
234       tracks_.GetMarkersInFrame(clip, frame, &this_frame_markers);
235       LG << "Clip " << clip << ", frame " << frame << " have "
236          << this_frame_markers.size();
237       if (this_frame_markers.size() < options.min_num_features) {
238         DetectFeaturesInFrame(clip, frame);
239         this_frame_markers.clear();
240         tracks_.GetMarkersInFrame(clip, frame, &this_frame_markers);
241         LG << "... detected " << this_frame_markers.size() << " features.";
242       }
243       if (previous_frame_markers.empty()) {
244         LG << "First frame; skipping tracking stage.";
245         previous_frame_markers.swap(this_frame_markers);
246         continue;
247       }
248       // Second, find tracks that should get tracked forward into this frame.
249       // To avoid tracking markers that are already tracked to this frame, make
250       // a sorted set of the tracks that exist in the last frame.
251       vector<int> tracks_in_this_frame;
252       for (int i = 0; i < this_frame_markers.size(); ++i) {
253         tracks_in_this_frame.push_back(this_frame_markers[i].track);
254       }
255       std::sort(tracks_in_this_frame.begin(),
256                 tracks_in_this_frame.end());
257
258       // Find tracks in the previous frame that are not in this one.
259       vector<Marker*> previous_frame_markers_to_track;
260       int num_skipped = 0;
261       for (int i = 0; i < previous_frame_markers.size(); ++i) {
262         if (std::binary_search(tracks_in_this_frame.begin(),
263                                tracks_in_this_frame.end(),
264                                previous_frame_markers[i].track)) {
265           num_skipped++;
266         } else {
267           previous_frame_markers_to_track.push_back(&previous_frame_markers[i]);
268         }
269       }
270
271       // Finally track the markers from the last frame into this one.
272       // TODO(keir): Use OMP.
273       for (int i = 0; i < previous_frame_markers_to_track.size(); ++i) {
274         Marker this_frame_marker = *previous_frame_markers_to_track[i];
275         this_frame_marker.frame = frame;
276         LG << "Tracking: " << this_frame_marker;
277         TrackRegionResult result;
278         TrackMarker(&this_frame_marker, &result);
279         if (result.is_usable()) {
280           LG << "Success: " << this_frame_marker;
281           AddMarker(this_frame_marker);
282           this_frame_markers.push_back(this_frame_marker);
283         } else {
284           LG << "Failed to track: " << this_frame_marker;
285         }
286       }
287       // Put the markers from this frame
288       previous_frame_markers.swap(this_frame_markers);
289     }
290   }
291 }
292
293 }  // namespace mv