main.cpp 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180
  1. //
  2. // Created by wrede on 19.04.16.
  3. //
  4. #include "../core/Definitions.h"
  5. #include "../core/DetectionSequence.h"
  6. #include "../util/FileIO.h"
  7. #include "../util/Parser.h"
  8. #include "../algo/TwoStage.h"
  9. #include "../visual/Visualizer.h"
  10. #include "../util/Logger.h"
  11. #include "../core/ObjectDataAngular.h"
  12. #include <boost/program_options.hpp>
  13. void ReadInput(const std::string& input_file, core::DetectionSequence& sequence,
  14. double temporal_weight, double spatial_weight, double angular_weight)
  15. {
  16. util::Logger::LogInfo("Reading input");
  17. core::Vector3d values;
  18. util::FileIO::ReadCSV(values, input_file);
  19. util::Parser::ParseObjectDataAngular(values, sequence,
  20. temporal_weight,
  21. spatial_weight,
  22. angular_weight);
  23. if (util::Logger::IsDebugEnabled())
  24. {
  25. size_t sequence_object_count = 0;
  26. for (size_t i = 0; i < sequence.GetFrameCount(); i++)
  27. {
  28. sequence_object_count += sequence.GetObjectCount(i);
  29. }
  30. util::Logger::LogDebug("sequence object count " + std::to_string(sequence_object_count));
  31. }
  32. }
  33. struct
  34. {
  35. size_t max_frame_skip;
  36. size_t max_tracklet_count;
  37. double penalty_value;
  38. } two_stage_params;
  39. void RunTwoStage(core::DetectionSequence& sequence, const std::string& output_file,
  40. const std::string& images_folder, bool display)
  41. {
  42. util::Logger::LogInfo("Running two-stage");
  43. algo::TwoStage two_stage(two_stage_params.max_frame_skip,
  44. two_stage_params.penalty_value,
  45. two_stage_params.max_tracklet_count);
  46. // Running the two stage graph algorithm
  47. algo::DirectedGraph obj_graph;
  48. two_stage.CreateObjectGraph(obj_graph, sequence);
  49. algo::DirectedGraph tlt_graph_1;
  50. two_stage.CreateTrackletGraph(obj_graph, tlt_graph_1, sequence.GetFrameCount());
  51. algo::DirectedGraph tlt_graph_2;
  52. two_stage.CreateTrackletGraph(tlt_graph_1, tlt_graph_2, sequence.GetFrameCount());
  53. std::vector<core::TrackletPtr> tracks;
  54. two_stage.ExtractTracks(tlt_graph_2, 1, tracks);
  55. // Interpolate tracks
  56. for (auto track : tracks)
  57. {
  58. track->InterpolateMissingFrames();
  59. }
  60. // Display the tracking data
  61. if (display)
  62. {
  63. util::Logger::LogInfo("Displaying data");
  64. visual::Visualizer vis;
  65. vis.Display(tracks, images_folder);
  66. }
  67. util::Logger::LogInfo("Finished");
  68. }
  69. void Run(int argc, char** argv)
  70. {
  71. // Algorithm independent values
  72. std::string input_file, output_file, images_folder, algorithm;
  73. bool display;
  74. // Input dependent variables
  75. double temporal_weight, spatial_weight, angular_weight;
  76. boost::program_options::options_description opts("Allowed options");
  77. opts.add_options()
  78. ("help",
  79. "produce help message")
  80. ("info",
  81. "if the program should show progress information")
  82. ("debug",
  83. "if the program should show debug messages")
  84. ("input-file,i",
  85. boost::program_options::value<std::string>(&input_file),
  86. "set detections file path")
  87. ("output-file,o",
  88. boost::program_options::value<std::string>(&output_file),
  89. "set the output file path")
  90. ("algorithm,a",
  91. boost::program_options::value<std::string>(&algorithm),
  92. "set the algorithm to use, current viable options: two-stage")
  93. ("display",
  94. "if a window with the images and the detected tracks should be opened")
  95. ("images-folder,f",
  96. boost::program_options::value<std::string>(&images_folder),
  97. "set images folder path")
  98. ("max-frame-skip",
  99. boost::program_options::value<size_t>(&two_stage_params.max_frame_skip)->default_value(1),
  100. "(two stage) set the maximum number of frames a track can skip between two detections")
  101. ("max-tracklet-count",
  102. boost::program_options::value<size_t>(&two_stage_params.max_tracklet_count)->default_value(1),
  103. "(two stage) set the maximum number of tracklets to be extracted")
  104. ("penalty-value",
  105. boost::program_options::value<double>(&two_stage_params.penalty_value)->default_value(0.0),
  106. "(two stage) set the penalty value for edges from and to source and sink")
  107. ("temporal-weight",
  108. boost::program_options::value<double>(&temporal_weight)->default_value(0.3),
  109. "temporal weight for difference calculations between two detections")
  110. ("spatial-weight",
  111. boost::program_options::value<double>(&spatial_weight)->default_value(0.3),
  112. "spatial weight for difference calculations between two detections")
  113. ("angular-weight",
  114. boost::program_options::value<double>(&angular_weight)->default_value(0.3),
  115. "angular weight for difference calculations between two detections");
  116. boost::program_options::variables_map opt_var_map;
  117. #pragma clang diagnostic push
  118. #pragma clang diagnostic ignored "-Wincompatible-pointer-types-discards-qualifiers"
  119. boost::program_options::store(boost::program_options::parse_command_line(argc, argv, opts), opt_var_map);
  120. #pragma clang diagnostic pop
  121. boost::program_options::notify(opt_var_map);
  122. if (opt_var_map.count("help") != 0)
  123. {
  124. std::cout << opts << std::endl;
  125. exit(0);
  126. }
  127. if (opt_var_map.count("info") != 0)
  128. {
  129. util::Logger::SetInfo(true);
  130. util::Logger::LogInfo("Enabled");
  131. }
  132. if (opt_var_map.count("debug") != 0)
  133. {
  134. util::Logger::SetDebug(true);
  135. util::Logger::LogDebug("Enabled");
  136. }
  137. display = opt_var_map.count("display") != 0;
  138. core::DetectionSequence sequence;
  139. ReadInput(input_file, sequence, temporal_weight, spatial_weight, angular_weight);
  140. if (algorithm == "two-stage")
  141. {
  142. RunTwoStage(sequence, output_file, images_folder, display);
  143. }
  144. else
  145. {
  146. std::cout << opts << std::endl;
  147. exit(0);
  148. }
  149. }
  150. int main(int argc, char** argv)
  151. {
  152. Run(argc, argv);
  153. //TestTracklet();
  154. return 0;
  155. }