README: mention some more video where the effect is used
[experiments/opencv_trail_effect.git] / 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