NStage.cpp 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299
  1. //
  2. // Created by wrede on 25.04.16.
  3. //
  4. #include "NStage.h"
  5. #include "../util/Logger.h"
  6. #include <boost/graph/dijkstra_shortest_paths.hpp>
  7. namespace algo
  8. {
  9. NStage::NStage(std::vector<size_t> max_frame_skip,
  10. std::vector<double> penalty_value,
  11. std::vector<size_t> max_tracklet_count,
  12. double edge_weight_threshold,
  13. std::vector<std::unordered_map<std::string, double>> constraints)
  14. {
  15. max_frame_skips_ = max_frame_skip;
  16. penalty_values_ = penalty_value;
  17. max_tracklet_counts_ = max_tracklet_count;
  18. iterations_ = std::min(max_tracklet_count.size(), penalty_value.size());
  19. edge_weight_threshold_ = edge_weight_threshold;
  20. constraints_ = constraints;
  21. }
  22. void NStage::CreateObjectGraph(DirectedGraph & graph,
  23. core::DetectionSequence & detections)
  24. {
  25. util::Logger::LogInfo("Creating object graph");
  26. std::vector<std::vector<Vertex>> layers;
  27. // Add source as the vertex with the lowest index
  28. Vertex source = boost::add_vertex(core::ObjectDataPtr(new core::ObjectData()), graph);
  29. // Add vertices from detection sequence to directed graph
  30. // Save the vertices which are in one frame/layer for later use to
  31. // link easily between vertices in adjacent frames/layers
  32. for (size_t i = 0; i < detections.GetFrameCount(); ++i)
  33. {
  34. std::vector<Vertex> layer;
  35. for (size_t j = 0; j < detections.GetObjectCount(i); ++j)
  36. {
  37. Vertex v = boost::add_vertex(detections.GetObject(i, j), graph);
  38. layer.push_back(v);
  39. }
  40. layers.push_back(layer);
  41. }
  42. // Add sink as the vertex with the highest index
  43. Vertex sink = boost::add_vertex(core::ObjectDataPtr(new core::ObjectData()), graph);
  44. // Vertex objects
  45. VertexValueMap values = boost::get(boost::vertex_name, graph);
  46. // Create edges
  47. for (size_t i = 0; i < layers.size(); ++i)
  48. {
  49. // For each edge in this frame/layer
  50. for (size_t j = 0; j < layers[i].size(); ++j)
  51. {
  52. Vertex u = layers[i][j];
  53. // For each next frame/layer until maxFrameSkip or end
  54. for (size_t k = 1; k <= max_frame_skips_[0] && i + k < layers.size(); ++k)
  55. {
  56. // To every edge in the next frame/layer
  57. for (size_t l = 0; l < layers[i + k].size(); ++l)
  58. {
  59. Vertex v = layers[i + k][l];
  60. // Only create the edge if the constraints are assured
  61. if (values[u]->IsWithinConstraints(values[v], constraints_[0]))
  62. {
  63. double weight = values[u]->CompareTo(values[v]);
  64. if (weight < edge_weight_threshold_)
  65. {
  66. boost::add_edge(u, v, weight, graph);
  67. }
  68. }
  69. }
  70. }
  71. // From source to vertex and from vertex to sink
  72. boost::add_edge(source, u,
  73. (i + 1) * penalty_values_[0],
  74. graph);
  75. boost::add_edge(u, sink,
  76. (layers.size() - i) * penalty_values_[0],
  77. graph);
  78. }
  79. }
  80. util::Logger::LogDebug("vertex count " + std::to_string(boost::num_vertices(graph)));
  81. util::Logger::LogDebug("edge count " + std::to_string(boost::num_edges(graph)));
  82. }
  83. void NStage::CreateTrackletGraph(DirectedGraph& obj_graph, DirectedGraph& tlt_graph,
  84. size_t frame_count, size_t iteration)
  85. {
  86. util::Logger::LogInfo("Creating tracklet graph");
  87. // Add source to tracklet graph
  88. Vertex tlt_src = boost::add_vertex(
  89. core::ObjectDataPtr(new core::ObjectData()), tlt_graph);
  90. // Prepare variables for dijkstra
  91. size_t obj_graph_size = boost::num_vertices(obj_graph);
  92. std::vector<Vertex> obj_pred_list(obj_graph_size);
  93. std::vector<double> obj_dist_list(obj_graph_size);
  94. VertexIndexMap obj_indices = boost::get(boost::vertex_index, obj_graph);
  95. VertexValueMap obj_values = boost::get(boost::vertex_name, obj_graph);
  96. PredecessorMap obj_pred_map(&obj_pred_list[0], obj_indices);
  97. DistanceMap obj_dist_map(&obj_dist_list[0], obj_indices);
  98. // Source and sink of the object graph
  99. Vertex obj_src = obj_indices[0];
  100. Vertex obj_snk = obj_indices[obj_graph_size - 1];
  101. //TODO experimental
  102. EdgeWeightMap weight_map = boost::get(boost::edge_weight, obj_graph);
  103. // Iteratively run dijkstra to extract tracklets
  104. for (size_t i = 0; i != max_tracklet_counts_[iteration]; ++i)
  105. {
  106. util::Logger::LogDebug("tracklet iteration: " + std::to_string(i));
  107. boost::dijkstra_shortest_paths(obj_graph, obj_src,
  108. boost::predecessor_map(obj_pred_map)
  109. .distance_map(obj_dist_map));
  110. // No path from source to sink could be found
  111. if (obj_dist_map[obj_snk] == std::numeric_limits<double>::max())
  112. {
  113. break;
  114. }
  115. // Create the tracklet
  116. core::TrackletPtr tracklet(new core::Tracklet);
  117. for (Vertex u = obj_pred_map[obj_snk], v = obj_snk; u != v; v = u, u = obj_pred_map[v])
  118. {
  119. tracklet->AddPathObject(obj_values[u]);
  120. // Leave source and sink untouched
  121. if (!obj_values[u]->IsVirtual())
  122. {
  123. //TODO original
  124. // Remove the path by setting all used edges to a weight of
  125. // infinity
  126. // std::pair<DirectedGraph::out_edge_iterator,
  127. // DirectedGraph::out_edge_iterator> edge_iter = boost::out_edges(u, obj_graph);
  128. //
  129. // for (DirectedGraph::out_edge_iterator iter = edge_iter.first;
  130. // iter != edge_iter.second;
  131. // ++iter)
  132. // {
  133. // boost::get(boost::edge_weight, obj_graph, *iter)
  134. // = std::numeric_limits<double>::infinity();
  135. // }
  136. //TODO experimental
  137. OutEdgeIter oei, oei_end;
  138. for (boost::tie(oei, oei_end) = boost::out_edges(u, obj_graph);
  139. oei != oei_end;
  140. ++oei)
  141. {
  142. weight_map[*oei] = std::numeric_limits<double>::infinity();
  143. }
  144. }
  145. }
  146. core::ObjectDataPtr tracklet_base = tracklet;
  147. // Add tracklet into tracklet graph
  148. boost::add_vertex(tracklet_base, tlt_graph);
  149. }
  150. // Add sink to tracklet graph
  151. Vertex tlt_snk = boost::add_vertex(core::ObjectDataPtr(new core::ObjectData()),
  152. tlt_graph);
  153. util::Logger::LogDebug("adding edges");
  154. // Create edges
  155. size_t tlt_graph_size = boost::num_vertices(tlt_graph);
  156. VertexIndexMap tlt_indices = boost::get(boost::vertex_index, tlt_graph);
  157. VertexValueMap tlt_values = boost::get(boost::vertex_name, tlt_graph);
  158. // For every tracklet but source and sink
  159. for (size_t i = 1; i < tlt_graph_size - 1; ++i)
  160. {
  161. Vertex u = tlt_indices[i];
  162. core::TrackletPtr u_ptr = std::static_pointer_cast<core::Tracklet>(tlt_values[u]);
  163. size_t u_first_frame = u_ptr->GetFirstFrameIndex();
  164. size_t u_last_frame = u_ptr->GetLastFrameIndex();
  165. // Create edges between tracklets
  166. for (size_t j = 1; j < tlt_graph_size - 1; ++j)
  167. {
  168. if (i != j)
  169. {
  170. Vertex v = tlt_indices[j];
  171. core::TrackletPtr v_ptr =
  172. std::static_pointer_cast<core::Tracklet>(tlt_values[v]);
  173. size_t v_first_frame = v_ptr->GetFirstFrameIndex();
  174. // Link only tracklets that are in temporal order
  175. if (u_last_frame < v_first_frame &&
  176. (v_first_frame - u_last_frame < max_frame_skips_[iteration]))
  177. {
  178. // Only create the edge if the constraints are assured
  179. if (u_ptr->IsWithinConstraints(v_ptr, constraints_[iteration]))
  180. {
  181. boost::add_edge(u, v,
  182. tlt_values[u]->CompareTo(tlt_values[v]),
  183. tlt_graph);
  184. }
  185. }
  186. }
  187. }
  188. // From source
  189. boost::add_edge(tlt_src, u,
  190. (u_first_frame + 1) * penalty_values_[iteration],
  191. tlt_graph);
  192. // To sink
  193. boost::add_edge(u, tlt_snk,
  194. (frame_count - u_last_frame) * penalty_values_[iteration],
  195. tlt_graph);
  196. }
  197. util::Logger::LogDebug("vertex count " + std::to_string(boost::num_vertices(tlt_graph)));
  198. util::Logger::LogDebug("edge count " + std::to_string(boost::num_edges(tlt_graph)));
  199. }
  200. void NStage::ExtractTracks(DirectedGraph& tlt_graph, size_t depth,
  201. std::vector<core::TrackletPtr>& tracks)
  202. {
  203. util::Logger::LogInfo("Extracting tracks");
  204. VertexValueMap tlt_values = boost::get(boost::vertex_name, tlt_graph);
  205. for (size_t i = 0; i < boost::num_vertices(tlt_graph); ++i)
  206. {
  207. core::ObjectDataPtr obj = tlt_values[i];
  208. if (!obj->IsVirtual())
  209. {
  210. core::TrackletPtr tlt = std::static_pointer_cast<core::Tracklet>(obj);
  211. for (size_t j = 0; j < depth; j++)
  212. {
  213. tlt->Flatten();
  214. }
  215. tracks.push_back(tlt);
  216. }
  217. }
  218. util::Logger::LogDebug("track count " + std::to_string(tracks.size()));
  219. }
  220. void NStage::Run(core::DetectionSequence & sequence,
  221. std::vector<core::TrackletPtr>& tracks)
  222. {
  223. // Running the two stage graph algorithm
  224. DirectedGraph obj_graph;
  225. CreateObjectGraph(obj_graph, sequence);
  226. // Run the tracklet creation at least once
  227. DirectedGraph tlt_graph_1, tlt_graph_2;
  228. CreateTrackletGraph(obj_graph, tlt_graph_1, sequence.GetFrameCount(), 0);
  229. // Run the tracklet creation iteratively
  230. for (size_t i = 1; i < iterations_; ++i)
  231. {
  232. if (i % 2 == 0)
  233. {
  234. CreateTrackletGraph(tlt_graph_2, tlt_graph_1, sequence.GetFrameCount(), i);
  235. }
  236. else
  237. {
  238. CreateTrackletGraph(tlt_graph_1, tlt_graph_2, sequence.GetFrameCount(), i);
  239. }
  240. }
  241. // Extract tracklets and flatten tracklets
  242. if (iterations_ % 2 == 0)
  243. {
  244. ExtractTracks(tlt_graph_2, iterations_ - 1, tracks);
  245. }
  246. else
  247. {
  248. ExtractTracks(tlt_graph_1, iterations_ - 1, tracks);
  249. }
  250. }
  251. }