Segmentation.hpp

   1 /*
   2  * openvc_trail_effect - experiments about video trail effects
   3  *
   4  * Copyright (C) 2015  Antonio Ospite <ao2@ao2.it>
   5  *
   6  * This program is free software: you can redistribute it and/or modify
   7  * it under the terms of the GNU General Public License as published by
   8  * the Free Software Foundation, either version 2 of the License, or
   9  * (at your option) any later version.
  10  *
  11  * This program is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14  * GNU General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  18  */
  19
  20 #ifndef SEGMENTATION_HPP
  21 #define SEGMENTATION_HPP
  22
  23 #include <opencv2/opencv.hpp>
  24
  25 class Segmentation {
  26 public:
  27         virtual cv::Mat getForegroundMask(const cv::Mat& frame) = 0;
  28         virtual ~Segmentation() {}
  29 };
  30
  31 class DummySegmentation : public Segmentation {
  32 public:
  33         cv::Mat getForegroundMask(const cv::Mat& frame)
  34         {
  35                 return cv::Mat(frame.size(), CV_8UC1, cv::Scalar(1));
  36         }
  37 };
  38
  39 class ThresholdSegmentation : public Segmentation {
  40 private:
  41         int threshold;
  42
  43 public:
  44         ThresholdSegmentation(int threshold_value)
  45         {
  46                 threshold = threshold_value;
  47         }
  48
  49         cv::Mat getForegroundMask(const cv::Mat& frame)
  50         {
  51                 cv::Mat gray_frame;
  52                 cv::Mat frame_mask;
  53
  54                 cvtColor(frame, gray_frame, CV_RGB2GRAY);
  55                 cv::threshold(gray_frame, frame_mask, threshold, 255, CV_THRESH_TOZERO);
  56
  57                 return frame_mask;
  58         }
  59 };
  60
  61 /*
  62  * Some background on... background subtraction can be found here:
  63  * http://docs.opencv.org/master/d1/dc5/tutorial_background_subtraction.html
  64  */
  65
  66 class MOG2Segmentation : public Segmentation {
  67 public:
  68         MOG2Segmentation(cv::VideoCapture& inputVideo, int learning_frames)
  69         {
  70                 cv::Mat background;
  71                 cv::Mat foreground_mask;
  72                 pMOG2 = cv::createBackgroundSubtractorMOG2();
  73
  74                 for (int i = 0; i < learning_frames; i++) {
  75                         inputVideo >> background;
  76                         pMOG2->apply(background, foreground_mask);
  77                 }
  78         }
  79
  80         cv::Mat getForegroundMask(const cv::Mat& frame)
  81         {
  82                 cv::Mat foreground_mask;
  83
  84                 pMOG2->apply(frame, foreground_mask, 0);
  85                 cv::erode(foreground_mask, foreground_mask, cv::Mat());
  86                 cv::dilate(foreground_mask, foreground_mask, cv::Mat());
  87                 cv::threshold(foreground_mask, foreground_mask, 0, 255, CV_THRESH_OTSU);
  88                 cv::medianBlur(foreground_mask, foreground_mask, 9);
  89
  90                 return foreground_mask;
  91         }
  92
  93         void getBackgroundImage(const cv::Mat& background)
  94         {
  95                 pMOG2->getBackgroundImage(background);
  96         }
  97
  98 private:
  99         cv::Ptr<cv::BackgroundSubtractor> pMOG2;
 100 };
 101
 102 #endif // SEGMENTATION_HPP