Berclaz.cpp 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213
  1. //
  2. // Created by wrede on 02.06.16.
  3. //
  4. #include "Berclaz.h"
  5. #include "../util/Parser.h"
  6. #include "../util/Logger.h"
  7. #include "KShortestPaths.h"
  8. namespace algo
  9. {
  10. Berclaz::Berclaz(int h_res, int v_res, int vicinity_size)
  11. {
  12. h_res_ = h_res;
  13. v_res_ = v_res;
  14. vicinity_size_ = vicinity_size;
  15. }
  16. void Berclaz::CreateGraph(DirectedGraph& graph, Vertex& source, Vertex& sink, util::Grid& grid)
  17. {
  18. util::Logger::LogDebug("add vertices");
  19. // Add grid vertices
  20. for (int z = 0; z < grid.GetDepthCount(); ++z)
  21. {
  22. for (int y = 0; y < grid.GetHeightCount(); ++y)
  23. {
  24. for (int x = 0; x < grid.GetWidthCount(); ++x)
  25. {
  26. boost::add_vertex(grid.GetValue(x, y, z), graph);
  27. }
  28. }
  29. }
  30. // Add source and sink vertex
  31. source = boost::add_vertex(core::ObjectDataPtr(new core::ObjectData()), graph);
  32. sink = boost::add_vertex(core::ObjectDataPtr(new core::ObjectData()), graph);
  33. util::Logger::LogDebug("add edges");
  34. // Iterate all vertices but source and sink
  35. VertexIndexMap vertices = boost::get(boost::vertex_index, graph);
  36. VertexValueMap values = boost::get(boost::vertex_name, graph);
  37. int layer_size = grid.GetWidthCount() * grid.GetHeightCount();
  38. for (int z = 0; z < grid.GetDepthCount(); ++z)
  39. {
  40. for (int y = 0; y < grid.GetHeightCount(); ++y)
  41. {
  42. for (int x = 0; x < grid.GetWidthCount(); ++x)
  43. {
  44. // First vertex index
  45. int vi = x + y * grid.GetWidthCount() + z * layer_size;
  46. // Get the score, clamp it, prevent division by zero and
  47. // logarithm of zero
  48. double score = values[vi]->GetDetectionScore();
  49. if (score > MAX_SCORE_VALUE)
  50. {
  51. score = MAX_SCORE_VALUE;
  52. }
  53. else if (score < MIN_SCORE_VALUE)
  54. {
  55. score = MIN_SCORE_VALUE;
  56. }
  57. // Calculate the edge weight
  58. double weight = -std::log(score / (1 - score));
  59. // Connect with the next frame only if there is a next frame
  60. if (z < grid.GetDepthCount() - 1)
  61. {
  62. // Iterate all nearby cells in the next frame
  63. for (int ny = std::max(0, y - vicinity_size_);
  64. ny < std::min(grid.GetHeightCount(), y + vicinity_size_ + 1);
  65. ++ny)
  66. {
  67. for (int nx = std::max(0, x - vicinity_size_);
  68. nx < std::min(grid.GetWidthCount(), x + vicinity_size_ + 1);
  69. ++nx)
  70. {
  71. // Second vertex index
  72. int vj = nx + ny * grid.GetWidthCount() + (z + 1) * layer_size;
  73. // Connect to nearby cells
  74. boost::add_edge(vertices[vi], vertices[vj], weight, graph);
  75. }
  76. }
  77. boost::add_edge(vertices[vi], sink, VIRTUAL_EDGE_WEIGHT, graph);
  78. }
  79. else
  80. {
  81. boost::add_edge(vertices[vi], sink, weight, graph);
  82. }
  83. // Connect with source and sink
  84. boost::add_edge(source, vertices[vi], VIRTUAL_EDGE_WEIGHT, graph);
  85. }
  86. }
  87. }
  88. util::Logger::LogDebug("vertex count " + std::to_string(boost::num_vertices(graph)));
  89. util::Logger::LogDebug("edge count " + std::to_string(boost::num_edges(graph)));
  90. }
  91. void Berclaz::ExtractTracks(DirectedGraph& graph, MultiPredecessorMap& map, Vertex origin,
  92. std::vector<core::TrackletPtr>& tracks)
  93. {
  94. VertexValueMap values = boost::get(boost::vertex_name, graph);
  95. // Move along all paths in reverse, starting at the origin
  96. for (Vertex first : map[origin])
  97. {
  98. core::TrackletPtr tracklet(new core::Tracklet());
  99. // The paths are node disjoint, so there should always be only one
  100. // node to proceed to
  101. for (Vertex u = first, v = (*map[u].begin());
  102. u != v; u = v, v = (*map[v].begin()))
  103. {
  104. tracklet->AddPathObject(values[u]);
  105. }
  106. tracks.push_back(tracklet);
  107. }
  108. }
  109. void Berclaz::Run(core::DetectionSequence& sequence,
  110. size_t batch_size, size_t max_track_count,
  111. std::vector<core::TrackletPtr> & tracks, util::Filter2D & filter)
  112. {
  113. size_t batch_count = 0;
  114. for (size_t i = sequence.GetFrameOffset(); i < sequence.GetFrameCount(); i += batch_size)
  115. {
  116. util::Logger::LogDebug("batch offset: " + std::to_string(i));
  117. util::Grid grid = util::Parser::ParseGrid(sequence, i, i + batch_size,
  118. 0.0, 1.0, h_res_, 0.0, 1.0, v_res_);
  119. grid.Convolve2D(filter);
  120. util::Logger::LogDebug("create graph");
  121. DirectedGraph graph;
  122. Vertex source, sink;
  123. CreateGraph(graph, source, sink, grid);
  124. util::Logger::LogDebug("run ksp");
  125. KShortestPaths ksp(graph, source, sink);
  126. ksp.Run(max_track_count);
  127. util::Logger::LogDebug("get paths");
  128. std::vector<std::vector<Vertex>> paths;
  129. ksp.GetPaths(paths);
  130. util::Logger::LogDebug("extract tracks");
  131. VertexValueMap values = boost::get(boost::vertex_name, graph);
  132. util::Logger::LogDebug("");
  133. for (auto p : paths)
  134. {
  135. std::string path_string = "";
  136. core::TrackletPtr tlt(new core::Tracklet());
  137. for (auto v : p)
  138. {
  139. path_string += std::to_string(v) + " ";
  140. tlt->AddPathObject(values[v]);
  141. }
  142. tracks.push_back(tlt);
  143. util::Logger::LogDebug(path_string);
  144. }
  145. util::Logger::LogDebug("");
  146. batch_count++;
  147. }
  148. // Only connect tracks if the sequence was split
  149. if (batch_count > 1)
  150. {
  151. //TODO find a better way to connect tracks (n-stage)
  152. util::Logger::LogDebug("connect tracks");
  153. ConnectTracks(tracks);
  154. }
  155. }
  156. void Berclaz::ConnectTracks(std::vector<core::TrackletPtr>& tracks)
  157. {
  158. for (size_t i = 0; i < tracks.size(); ++i)
  159. {
  160. // find the best matching tracklet
  161. double best_value = std::numeric_limits<double>::max();
  162. size_t best_index = 0;
  163. for (size_t k = i + 1; k < tracks.size(); ++k)
  164. {
  165. if (tracks[i]->GetLastFrameIndex() < tracks[k]->GetFirstFrameIndex())
  166. {
  167. double value = tracks[i]->CompareTo(tracks[k]);
  168. if (value < best_value)
  169. {
  170. best_index = k;
  171. }
  172. }
  173. }
  174. // if a match was found
  175. if (best_index != 0)
  176. {
  177. // merge the two tracks
  178. tracks[i]->Combine(tracks[best_index]);
  179. tracks.erase(tracks.begin() + best_index);
  180. }
  181. }
  182. }
  183. }