Advertisement
dan-masek

Refactor of Crazed's frame processing into a class.

Jul 28th, 2017
222
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 3.59 KB | None | 0 0
  1. // ============================================================================
  2. class frame_processor
  3. {
  4. public:
  5.     typedef std::vector<cv::Point> contour_t;
  6.  
  7.     struct polygon_type
  8.     {
  9.         cv::Scalar color;
  10.         std::string name;
  11.     };
  12.  
  13.     struct polygon_info
  14.     {
  15.         polygon_type type;
  16.         contour_t points;
  17.         cv::Point centroid;
  18.     };
  19.  
  20.     typedef std::vector<polygon_info> polygon_infos_t;
  21.  
  22. public:
  23.     frame_processor();
  24.  
  25.     polygon_infos_t process_frame(cv::Mat const& frame);
  26.  
  27. private:
  28.     typedef std::map<size_t, polygon_type> polygon_types_t;
  29.  
  30.     // void perform_morph_ops(cv::Mat &thresh) const; // Unused
  31.  
  32.     bool process_contour(contour_t const& contour, polygon_info& info);
  33.  
  34. private:
  35.     polygon_types_t polygon_types_;
  36.  
  37.     double threshold_level_;
  38.     double min_area_;
  39.  
  40.     cv::Mat erode_element_;
  41.     cv::Mat dilate_element_;
  42.  
  43.     cv::Mat grayscale_buffer_;
  44. };
  45. // ============================================================================
  46. frame_processor::frame_processor()
  47.     : threshold_level_(128)
  48.     , min_area_(100)
  49. {
  50.     erode_element_ = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(10, 10));
  51.     dilate_element_ = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(10, 10));
  52.  
  53.     // Initialize supported polygon types...
  54.     polygon_types_.emplace(3, polygon_type{ cv::Scalar(0, 0, 255), "Triangle" });
  55.     polygon_types_.emplace(4, polygon_type{ cv::Scalar(0, 255, 0), "Quadrilateral" });
  56.     polygon_types_.emplace(7, polygon_type{ cv::Scalar(255, 0, 0), "Heptagon" });
  57.     polygon_types_.emplace(10, polygon_type{ cv::Scalar(127, 127, 0), "Decagon" });
  58. }
  59. // ----------------------------------------------------------------------------
  60. /* Unused
  61. void frame_processor::perform_morph_ops(cv::Mat &thresh) const
  62. {
  63.     cv::erode(thresh, thresh, erode_element_, cv::Point(-1, -1), 2);
  64.     cv::dilate(thresh, thresh, dilate_element_, cv::Point(-1, -1), 2);
  65. }
  66. */
  67. // ----------------------------------------------------------------------------
  68. bool frame_processor::process_contour(contour_t const& contour, polygon_info& info)
  69. {
  70.     if (contourArea(contour) < min_area_) {
  71.         return false;
  72.     }
  73.    
  74.     cv::Moments moment = cv::moments(cv::Mat(contour));
  75.     int x1(static_cast<int>(std::round(moment.m10 / moment.m00)));
  76.     int y1(static_cast<int>(std::round(moment.m01 / moment.m00)));
  77.  
  78.     polygon_types_t::const_iterator it(polygon_types_.find(contour.size()));
  79.     if (it == polygon_types_.end()) {
  80.         return false;
  81.     }
  82.  
  83.     info.type = it->second;
  84.     info.points = contour;
  85.     info.centroid = cv::Point(x1, y1);
  86.  
  87.     return true;
  88. }
  89. // ----------------------------------------------------------------------------
  90. frame_processor::polygon_infos_t frame_processor::process_frame(cv::Mat const& frame)
  91. {
  92.     cv::cvtColor(frame, grayscale_buffer_, cv::COLOR_BGR2GRAY);
  93.  
  94.     cv::threshold(grayscale_buffer_, grayscale_buffer_, threshold_level_
  95.         , 255, cv::THRESH_BINARY);
  96.  
  97.     std::vector<contour_t> contours;
  98.     std::vector<cv::Vec4i> hierarchy;
  99.     cv::findContours(grayscale_buffer_, contours, hierarchy
  100.         , CV_RETR_TREE, CV_CHAIN_APPROX_SIMPLE);
  101.  
  102.     polygon_infos_t result;
  103.     for (contour_t const& contour : contours) {
  104.         std::vector<cv::Point> approx_contour;
  105.         cv::approxPolyDP(cv::Mat(contour), approx_contour, 3, true);
  106.         polygon_info info;
  107.         if (process_contour(approx_contour, info)) {
  108.             result.push_back(info);
  109.         }
  110.     }
  111.     return result;
  112. }
  113. // ============================================================================
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement