converted some more samples to C++
[profile/ivi/opencv.git] / samples / c / bgfg_codebook.cpp
1 // Background average sample code done with averages and done with codebooks
2 // (adapted from the OpenCV book sample)
3 // 
4 // NOTE: To get the keyboard to work, you *have* to have one of the video windows be active
5 //       and NOT the consule window.
6 //
7 // Gary Bradski Oct 3, 2008.
8 // 
9 /* *************** License:**************************
10    Oct. 3, 2008
11    Right to use this code in any way you want without warrenty, support or any guarentee of it working.
12
13    BOOK: It would be nice if you cited it:
14    Learning OpenCV: Computer Vision with the OpenCV Library
15      by Gary Bradski and Adrian Kaehler
16      Published by O'Reilly Media, October 3, 2008
17  
18    AVAILABLE AT: 
19      http://www.amazon.com/Learning-OpenCV-Computer-Vision-Library/dp/0596516134
20      Or: http://oreilly.com/catalog/9780596516130/
21      ISBN-10: 0596516134 or: ISBN-13: 978-0596516130    
22 ************************************************** */
23
24 #include <stdio.h>
25 #include <stdlib.h>
26 #include <ctype.h>
27
28 #include <opencv2/video/background_segm.hpp>
29 #include <opencv2/imgproc/imgproc_c.h>
30 #include <opencv2/highgui/highgui.hpp>
31
32 //VARIABLES for CODEBOOK METHOD:
33 CvBGCodeBookModel* model = 0;
34 const int NCHANNELS = 3;
35 bool ch[NCHANNELS]={true,true,true}; // This sets what channels should be adjusted for background bounds
36
37 void help(void)
38 {
39     printf("\nLearn background and find foreground using simple average and average difference learning method:\n"
40         "\nUSAGE:\nbgfg_codebook [--nframes=300] [movie filename, else from camera]\n"
41         "***Keep the focus on the video windows, NOT the consol***\n\n"
42         "INTERACTIVE PARAMETERS:\n"
43         "\tESC,q,Q  - quit the program\n"
44         "\th    - print this help\n"
45         "\tp    - pause toggle\n"
46         "\ts    - single step\n"
47         "\tr    - run mode (single step off)\n"
48         "=== AVG PARAMS ===\n"
49         "\t-    - bump high threshold UP by 0.25\n"
50         "\t=    - bump high threshold DOWN by 0.25\n"
51         "\t[    - bump low threshold UP by 0.25\n"
52         "\t]    - bump low threshold DOWN by 0.25\n"
53         "=== CODEBOOK PARAMS ===\n"
54         "\ty,u,v- only adjust channel 0(y) or 1(u) or 2(v) respectively\n"
55         "\ta    - adjust all 3 channels at once\n"
56         "\tb    - adjust both 2 and 3 at once\n"
57         "\ti,o  - bump upper threshold up,down by 1\n"
58         "\tk,l  - bump lower threshold up,down by 1\n"
59         "\tSPACE - reset the model\n"
60         );
61 }
62
63 //
64 //USAGE:  ch9_background startFrameCollection# endFrameCollection# [movie filename, else from camera]
65 //If from AVI, then optionally add HighAvg, LowAvg, HighCB_Y LowCB_Y HighCB_U LowCB_U HighCB_V LowCB_V
66 //
67 int main(int argc, char** argv)
68 {
69     const char* filename = 0;
70     IplImage* rawImage = 0, *yuvImage = 0; //yuvImage is for codebook method
71     IplImage *ImaskCodeBook = 0,*ImaskCodeBookCC = 0;
72     CvCapture* capture = 0;
73
74     int c, n, nframes = 0;
75     int nframesToLearnBG = 300;
76
77     model = cvCreateBGCodeBookModel();
78     
79     //Set color thresholds to default values
80     model->modMin[0] = 3;
81     model->modMin[1] = model->modMin[2] = 3;
82     model->modMax[0] = 10;
83     model->modMax[1] = model->modMax[2] = 10;
84     model->cbBounds[0] = model->cbBounds[1] = model->cbBounds[2] = 10;
85
86     bool pause = false;
87     bool singlestep = false;
88
89     for( n = 1; n < argc; n++ )
90     {
91         static const char* nframesOpt = "--nframes=";
92         if( strncmp(argv[n], nframesOpt, strlen(nframesOpt))==0 )
93         {
94             if( sscanf(argv[n] + strlen(nframesOpt), "%d", &nframesToLearnBG) == 0 )
95             {
96                 help();
97                 return -1;
98             }
99         }
100         else
101             filename = argv[n];
102     }
103
104     if( !filename )
105     {
106         printf("Capture from camera\n");
107         capture = cvCaptureFromCAM( 0 );
108     }
109     else
110     {
111         printf("Capture from file %s\n",filename);
112         capture = cvCreateFileCapture( filename );
113     }
114
115     if( !capture )
116     {
117         printf( "Can not initialize video capturing\n\n" );
118         help();
119         return -1;
120     }
121
122     //MAIN PROCESSING LOOP:
123     for(;;)
124     {
125         if( !pause )
126         {
127             rawImage = cvQueryFrame( capture );
128             ++nframes;
129             if(!rawImage) 
130                 break;
131         }
132         if( singlestep )
133             pause = true;
134         
135         //First time:
136         if( nframes == 1 && rawImage )
137         {
138             // CODEBOOK METHOD ALLOCATION
139             yuvImage = cvCloneImage(rawImage);
140             ImaskCodeBook = cvCreateImage( cvGetSize(rawImage), IPL_DEPTH_8U, 1 );
141             ImaskCodeBookCC = cvCreateImage( cvGetSize(rawImage), IPL_DEPTH_8U, 1 );
142             cvSet(ImaskCodeBook,cvScalar(255));
143             
144             cvNamedWindow( "Raw", 1 );
145             cvNamedWindow( "ForegroundCodeBook",1);
146             cvNamedWindow( "CodeBook_ConnectComp",1);
147         }
148
149         // If we've got an rawImage and are good to go:                
150         if( rawImage )
151         {
152             cvCvtColor( rawImage, yuvImage, CV_BGR2YCrCb );//YUV For codebook method
153             //This is where we build our background model
154             if( !pause && nframes-1 < nframesToLearnBG  )
155                 cvBGCodeBookUpdate( model, yuvImage );
156
157             if( nframes-1 == nframesToLearnBG  )
158                 cvBGCodeBookClearStale( model, model->t/2 );
159             
160             //Find the foreground if any
161             if( nframes-1 >= nframesToLearnBG  )
162             {
163                 // Find foreground by codebook method
164                 cvBGCodeBookDiff( model, yuvImage, ImaskCodeBook );
165                 // This part just to visualize bounding boxes and centers if desired
166                 cvCopy(ImaskCodeBook,ImaskCodeBookCC);  
167                 cvSegmentFGMask( ImaskCodeBookCC );
168             }
169             //Display
170             cvShowImage( "Raw", rawImage );
171             cvShowImage( "ForegroundCodeBook",ImaskCodeBook);
172             cvShowImage( "CodeBook_ConnectComp",ImaskCodeBookCC);
173         }
174
175         // User input:
176         c = cvWaitKey(10)&0xFF;
177         c = tolower(c);
178         // End processing on ESC, q or Q
179         if(c == 27 || c == 'q')
180             break;
181         //Else check for user input
182         switch( c )
183         {
184         case 'h':
185             help();
186             break;
187         case 'p':
188             pause = !pause;
189             break;
190         case 's':
191             singlestep = !singlestep;
192             pause = false;
193             break;
194         case 'r':
195             pause = false;
196             singlestep = false;
197             break;
198         case ' ':
199             cvBGCodeBookClearStale( model, 0 );
200             nframes = 0;
201             break;
202             //CODEBOOK PARAMS
203         case 'y': case '0':
204         case 'u': case '1':
205         case 'v': case '2':
206         case 'a': case '3':
207         case 'b': 
208             ch[0] = c == 'y' || c == '0' || c == 'a' || c == '3';
209             ch[1] = c == 'u' || c == '1' || c == 'a' || c == '3' || c == 'b';
210             ch[2] = c == 'v' || c == '2' || c == 'a' || c == '3' || c == 'b';
211             printf("CodeBook YUV Channels active: %d, %d, %d\n", ch[0], ch[1], ch[2] );
212             break;
213         case 'i': //modify max classification bounds (max bound goes higher)
214         case 'o': //modify max classification bounds (max bound goes lower)
215         case 'k': //modify min classification bounds (min bound goes lower)
216         case 'l': //modify min classification bounds (min bound goes higher)
217             {
218             uchar* ptr = c == 'i' || c == 'o' ? model->modMax : model->modMin;
219             for(n=0; n<NCHANNELS; n++)
220             {
221                 if( ch[n] )
222                 {
223                     int v = ptr[n] + (c == 'i' || c == 'l' ? 1 : -1);
224                     ptr[n] = cv::saturate_cast<uchar>(v);
225                 }
226                 printf("%d,", ptr[n]);
227             }
228             printf(" CodeBook %s Side\n", c == 'i' || c == 'o' ? "High" : "Low" );
229             }
230             break;
231         }
232     }           
233     
234     cvReleaseCapture( &capture );
235     cvDestroyWindow( "Raw" );
236     cvDestroyWindow( "ForegroundCodeBook");
237     cvDestroyWindow( "CodeBook_ConnectComp");
238     return 0;
239 }