opencv_trail_effect.cpp: fix showing the background
[experiments/opencv_trail_effect.git] / opencv_trail_effect.cpp
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 #include <opencv2/opencv.hpp>
21 #include <unistd.h>
22
23 #include "Frame.hpp"
24 #include "Segmentation.hpp"
25 #include "Trail.hpp"
26
27 static void usage(const char *name)
28 {
29         std::cout << "usage: " << name << " [OPTIONS]" << std::endl;
30         std::cout << "OPTIONS:" << std::endl;
31         std::cout << "\t-i <file>\tthe input file (if missing, a webcam will be tried)" << std::endl;
32         std::cout << "\t-o <file>\tthe optional output file" << std::endl;
33         std::cout << "\t-l <number>\tthe trail length in frames" << std::endl;
34         std::cout << "\t\t\tthe default is 25" << std::endl;
35         std::cout << "\t\t\tNOTES:" << std::endl;
36         std::cout << "\t\t\t  * a negative value means an 'infinite' trail" << std::endl;
37         std::cout << "\t-s <method>\tthe image segmentation method" << std::endl;
38         std::cout << "\t\t\tvalid values are:" << std::endl;
39         std::cout << "\t\t\t none, threshold, background" << std::endl;
40         std::cout << "\t\t\tthe default is 'background'" << std::endl;
41         std::cout << "\t\t\tNOTES:" << std::endl;
42         std::cout << "\t\t\t  * 'none' is only useful with '-d average'" << std::endl;
43         std::cout << "\t-b <number>\tthe number of initial frames for background learning," << std::endl;
44         std::cout << "\t\t\tthe default is 50" << std::endl;
45         std::cout << "\t\t\tNOTES:" << std::endl;
46         std::cout << "\t\t\t  * only useful with '-s background'" << std::endl;
47         std::cout << "\t-t <number>\tthe level for the threshold segmentation method," << std::endl;
48         std::cout << "\t\t\tthe default is 5" << std::endl;
49         std::cout << "\t\t\tNOTES:" << std::endl;
50         std::cout << "\t\t\t  * only useful with '-s threshold'" << std::endl;
51         std::cout << "\t-d <method>\tthe trail drawing method" << std::endl;
52         std::cout << "\t\t\tvalid values are:" << std::endl;
53         std::cout << "\t\t\t  copy, accumulate, fadecopy, fadeaccumulate, average" << std::endl;
54         std::cout << "\t\t\tthe default is 'copy'" << std::endl;
55         std::cout << "\t\t\tNOTES:" << std::endl;
56         std::cout << "\t\t\t  * 'copy' is useless with '-s none'" << std::endl;
57         std::cout << "\t\t\t  * the difference between 'fadecopy' and" << std::endl;
58         std::cout << "\t\t\t    'fadeaccumulate is cleared when using '-B'" << std::endl;
59         std::cout << "\t-r\t\treverse the trail drawing sequence" << std::endl;
60         std::cout << "\t-B\t\tshow the background behind the trail" << std::endl;
61         std::cout << "\t\t\tNOTES:" << std::endl;
62         std::cout << "\t\t\t  * only used with '-s background'" << std::endl;
63         std::cout << "\t-F\t\tredraw the current frame on top of the trail" << std::endl;
64         std::cout << "\t\t\tNOTES:" << std::endl;
65         std::cout << "\t\t\t  * noticeable with '-s average'" << std::endl;
66         std::cout << "\t\t\t  * noticeable with reverse faded trails" << std::endl;
67 }
68
69 int main(int argc, char *argv[])
70 {
71         int ret = 0;
72         int opt;
73
74         std::string *input_file = NULL;
75         std::string *output_file = NULL;
76         int trail_lenght = 25;
77         std::string *segmentation_method = new std::string("background");
78         int background_learn_frames = 50;
79         int threshold_level = 5;
80         std::string *drawing_method = new std::string("copy");
81         bool reverse_trail = false;
82         bool show_background = false;
83         bool redraw_current_frame = false;
84
85         while ((opt = getopt(argc, argv, "i:o:l:s:b:t:d:rBFh")) != -1) {
86                 switch (opt) {
87                 case 'i':
88                         input_file = new std::string(optarg);
89                         break;
90                 case 'o':
91                         output_file = new std::string(optarg);
92                         break;
93                 case 'l':
94                         trail_lenght = atoi(optarg);
95                         break;
96                 case 's':
97                         delete segmentation_method;
98                         segmentation_method = new std::string(optarg);
99                         break;
100                 case 'b':
101                         background_learn_frames = atoi(optarg);
102                         break;
103                 case 't':
104                         threshold_level = atoi(optarg);
105                         break;
106                 case 'd':
107                         delete drawing_method;
108                         drawing_method = new std::string(optarg);
109                         break;
110                 case 'r':
111                         reverse_trail = true;
112                         break;
113                 case 'B':
114                         show_background = true;
115                         break;
116                 case 'F':
117                         redraw_current_frame = true;
118                         break;
119                 case 'h':
120                         usage(argv[0]);
121                         return 0;
122                 default: /* '?' */
123                         usage(argv[0]);
124                         return -1;
125                 }
126         }
127
128         cv::VideoCapture inputVideo;
129         cv::VideoWriter outputVideo;
130         cv::Size frame_size;
131         cv::Mat input_frame;
132
133         if (input_file) {
134                 inputVideo.open(*input_file);
135         } else {
136                 inputVideo.open(0);
137         }
138
139         if (!inputVideo.isOpened()) {
140                 std::cerr  << "Could not open the input video." << std::endl;
141                 ret = -1;
142                 goto out;
143         }
144
145         frame_size = cv::Size((int) inputVideo.get(CV_CAP_PROP_FRAME_WIDTH),
146                               (int) inputVideo.get(CV_CAP_PROP_FRAME_HEIGHT));
147
148         if (output_file) {
149                 int fps = inputVideo.get(CV_CAP_PROP_FPS);
150                 if (fps < 0)
151                         fps = 25;
152
153                 outputVideo.open(*output_file, CV_FOURCC('M','J','P','G'), fps, frame_size, true);
154                 if (!outputVideo.isOpened()) {
155                         std::cerr  << "Could not open the output video for write." << std::endl;
156                         ret = -1;
157                         goto out;
158                 }
159         }
160
161         Trail *trail;
162         if (reverse_trail)
163                 trail = new ReverseTrail(trail_lenght, frame_size);
164         else
165                 trail = new ForwardTrail(trail_lenght, frame_size);
166
167         trail->setRedrawCurrentFrame(redraw_current_frame);
168
169         if (trail->setDrawingMethod(*drawing_method) < 0) {
170                 std::cerr  << "Invalid drawing method." << std::endl;
171                 ret = -1;
172                 goto out_delete_trail;
173         }
174
175         Segmentation *segmentation;
176         if (*segmentation_method == "background") {
177                 segmentation = new MOG2Segmentation(inputVideo, background_learn_frames);
178                 if (show_background) {
179                         cv::Mat background(frame_size, inputVideo.get(CV_CAP_PROP_FORMAT));
180
181                         ((MOG2Segmentation *)segmentation)->getBackgroundImage(background);
182                         trail->setBackground(background);
183                 }
184         } else if (*segmentation_method == "threshold") {
185                 segmentation = new ThresholdSegmentation(threshold_level);
186         } else if (*segmentation_method == "none") {
187                 segmentation = new DummySegmentation();
188         } else {
189                 std::cerr  << "Invalid segmentation method." << std::endl;
190                 ret = -1;
191                 goto out_delete_trail;
192         }
193
194         cv::namedWindow("Frame", CV_WINDOW_AUTOSIZE);
195
196         for (;;) {
197                 inputVideo >> input_frame;
198
199                 Frame *foreground = new Frame(input_frame,
200                                               segmentation->getForegroundMask(input_frame));
201                 trail->update(foreground);
202
203                 cv::Mat canvas = cv::Mat::zeros(input_frame.size(), input_frame.type());
204                 trail->draw(canvas);
205
206                 cv::imshow("Frame", canvas);
207                 if (cv::waitKeyEx(30) >= 0)
208                         break;
209
210                 if (outputVideo.isOpened())
211                         outputVideo << canvas;
212         }
213
214         cv::destroyWindow("Frame");
215
216         delete segmentation;
217
218 out_delete_trail:
219         delete trail;
220 out:
221         delete drawing_method;
222         delete segmentation_method;
223         delete output_file;
224         delete input_file;
225         return ret;
226 }