CasualGamer

OpenCV object/color detection

Oct 17th, 2020
716
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. //Color.h
  2. #pragma once
  3. enum Color {BLUE, RED, GREEN, YELLOW, PURPLE};
  4.  
  5. //Ball.h
  6. #pragma once
  7.  
  8. #include "Color.h"
  9. #include <opencv2/core/types.hpp>
  10.  
  11. class Ball
  12. {
  13. public:
  14.     Color color;
  15.     cv::Point position;
  16.     cv::Rect rect;
  17.  
  18.     Ball();
  19.     Ball(Color color, cv::Rect rect, int x, int y);
  20.  
  21. };
  22.  
  23. //Ball.cpp
  24. #include "Ball.h"
  25.  
  26. Ball::Ball(){
  27. }
  28.  
  29. Ball::Ball(Color color, cv::Rect rect, int x, int y){
  30.     this->color = color;
  31.     this->rect = rect;
  32.     this->position = cv::Point(x, y);
  33. }
  34.  
  35. //main.cpp
  36. #include <opencv2/highgui.hpp>
  37. #include <opencv2/imgproc/imgproc.hpp>
  38.  
  39. #include "Ball.h"
  40.  
  41. cv::Scalar yellowLow = cv::Scalar(25, 130, 180);
  42. cv::Scalar yellowHigh = cv::Scalar(45, 255, 255);
  43. cv::Scalar greenLow = cv::Scalar(46, 40, 40);
  44. cv::Scalar greenHigh = cv::Scalar(70, 255, 255);
  45. cv::Scalar blueLow = cv::Scalar(100, 150, 150);
  46. cv::Scalar blueHigh = cv::Scalar(140, 255, 255);
  47. cv::Scalar purpleLow = cv::Scalar(148, 117, 89);
  48. cv::Scalar purpleHigh = cv::Scalar(152, 255, 255);
  49. cv::Scalar redLow = cv::Scalar(170, 140, 160);
  50. cv::Scalar redHigh = cv::Scalar(180, 255, 255);
  51.  
  52. std::vector<Ball> balls;
  53.  
  54. void GetBalls(cv::Mat img, cv::Scalar low, cv::Scalar high, Color color) {
  55.     cv::Mat mask;
  56.     cv::inRange(img, low, high, mask);
  57.     std::vector<std::vector<cv::Point> > contours;
  58.     cv::findContours(mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
  59.  
  60.     for (size_t i = 0; i < contours.size(); i++)
  61.     {
  62.         cv::Rect boundRect = boundingRect(contours[i]);
  63.         if (boundRect.area() > 350 && (boundRect.width < 70 || boundRect.height < 70)) {
  64.             balls.emplace_back(color, boundRect, boundRect.x + boundRect.width / 2, boundRect.y + boundRect.height / 2);
  65.         }
  66.     }
  67. }
  68.  
  69. void drawBalls(cv::Mat background) {
  70.     for (size_t i = 0; i < balls.size(); i++) {
  71.         switch (balls[i].color) {
  72.         case RED:
  73.             rectangle(background, balls[i].rect.tl(), balls[i].rect.br(), CV_RGB(255, 0, 0), 2);
  74.             break;
  75.         case BLUE:
  76.             rectangle(background, balls[i].rect.tl(), balls[i].rect.br(), CV_RGB(0, 0, 255), 2);
  77.             break;
  78.         case GREEN:
  79.             rectangle(background, balls[i].rect.tl(), balls[i].rect.br(), CV_RGB(0, 255, 0), 2);
  80.             break;
  81.         case YELLOW:
  82.             rectangle(background, balls[i].rect.tl(), balls[i].rect.br(), CV_RGB(255, 255, 0), 2);
  83.             break;
  84.         case PURPLE:
  85.             rectangle(background, balls[i].rect.tl(), balls[i].rect.br(), CV_RGB(128, 0, 128), 2);
  86.             break;
  87.         }
  88.     }
  89. }
  90.  
  91.  
  92. int main(){
  93.     cv::Mat target = cv::imread("target.jpg"); //cv::IMREAD_COLOR
  94.  
  95.     cv::Mat background;
  96.     target.copyTo(background);
  97.  
  98.     cv::cvtColor(target, target, cv::COLOR_BGR2HSV);
  99.  
  100.     cv::rectangle(target, cv::Point(0, 0), cv::Point(640, 30), (0, 0, 0), cv::FILLED);
  101.  
  102.     GetBalls(target, yellowLow, yellowHigh, Color::YELLOW); //find yellow balls
  103.     GetBalls(target, blueLow, blueHigh, Color::BLUE); //find blue balls
  104.     GetBalls(target, redLow, redHigh, Color::RED); //find red balls
  105.     GetBalls(target, greenLow, greenHigh, Color::GREEN); //find green balls
  106.     GetBalls(target, purpleLow, purpleHigh, Color::PURPLE); // find purple balls
  107.     drawBalls(background);
  108.  
  109.     cv::imshow("contours", background);
  110.     cv::waitKey(0);
  111. }
RAW Paste Data