]> git.treefish.org Git - stickletrack.git/blob - tracking.cpp
windows stay closed and stickletrack exits properly if tracking screen is closes.
[stickletrack.git] / tracking.cpp
1 #include <stdio.h>
2 #include <iostream>
3
4 using namespace std;
5
6 #include "tracking.h"
7 #include "stickletrack.h"
8
9 #define SEPCOST_AREA      0
10 #define SEPCOST_CIRCUM    1
11 #define SEPCOST_SHORTAXIS 2
12 #define SEPCOST_LONGAXIS  3
13 #define SEPCOST_ANGLE     4
14 #define SEPCOST_DIST      5
15
16 struct fishblobb {
17   double x;
18   double y;
19   bool identified;
20   int contourid;
21   vector<Point> *points;
22   double longaxis;
23   double shortaxis;
24   double angle;
25   int area;
26   int circum;
27   Point2f head, tail;
28 };
29
30 static int enablecontours=0;
31 static Mat mContours;
32
33 struct tagcost {
34   int tagid;
35   double tagcost;
36 };
37
38 struct blobbcost {
39   int blobbid;
40   vector<tagcost> tagcosts;
41 };
42
43 bool sort_tags (tagcost i, tagcost j) {
44   return i.tagcost < j.tagcost;
45 }
46
47 bool sort_blobbs (blobbcost i, blobbcost j) {
48   return i.tagcosts[0].tagcost < j.tagcosts[0].tagcost;
49 }
50
51 Mat& tracking_getFrame() {
52   return mContours;
53 }
54
55 int tracking_showFrame() {
56   if ( ! isWindowClosed("stickletrack_tracking_screen") ) {
57     imshow("stickletrack_tracking_screen", mContours);
58     return 0;
59   }
60   else
61     return 1;
62 }
63
64 bool tracking_isEnabled () {
65   return enablecontours;
66 }
67
68 void tracking_init() {
69   namedWindow("stickletrack_tracking_prefs", CV_WINDOW_KEEPRATIO);
70   namedWindow("stickletrack_tracking_screen", CV_WINDOW_KEEPRATIO);
71
72   createTrackbar("Enable", "stickletrack_tracking_prefs", &enablecontours, 1, NULL, 0);
73   createTrackbar("Manyfish", "stickletrack_tracking_prefs", &Prefs.manyfish, 10, NULL, 0);
74   createTrackbar("min area", "stickletrack_tracking_prefs", &Prefs.contours_minarea, 100, &trackbarCallbackUpdateNormPrefs, 0);
75   createTrackbar("max area", "stickletrack_tracking_prefs", &Prefs.contours_maxarea, 100, &trackbarCallbackUpdateNormPrefs, 0);
76   createTrackbar("min circum", "stickletrack_tracking_prefs", &Prefs.contours_mincircum, 100, &trackbarCallbackUpdateNormPrefs, 0);
77   createTrackbar("max circum", "stickletrack_tracking_prefs", &Prefs.contours_maxcircum, 100, &trackbarCallbackUpdateNormPrefs, 0);
78   createTrackbar("min shortaxis", "stickletrack_tracking_prefs", &Prefs.contours_minshortaxis, 100, &trackbarCallbackUpdateNormPrefs, 0);
79   createTrackbar("max shortaxis", "stickletrack_tracking_prefs", &Prefs.contours_maxshortaxis, 100, &trackbarCallbackUpdateNormPrefs, 0);
80   createTrackbar("min longaxis", "stickletrack_tracking_prefs", &Prefs.contours_minlongaxis, 100, &trackbarCallbackUpdateNormPrefs, 0);
81   createTrackbar("max longaxis", "stickletrack_tracking_prefs", &Prefs.contours_maxlongaxis, 100, &trackbarCallbackUpdateNormPrefs, 0);
82   createTrackbar("max speed", "stickletrack_tracking_prefs", &Prefs.contours_maxspeed, 100, &trackbarCallbackUpdateNormPrefs, 0);
83   createTrackbar("max rotation speed", "stickletrack_tracking_prefs", &Prefs.contours_maxrot, 100, &trackbarCallbackUpdateNormPrefs, 0);
84
85   mContours = Mat::zeros(Props.height, Props.width, CV_8UC3);
86 }
87
88 double kitvalue (double min, double max) {
89   return min + pow( cos(frametime*10), 2)*(max-min);
90 }
91
92 void findFishBlobbs(vector<fishblobb>& fishblobbs, vector< vector<Point> >& contours, vector<Vec4i>& hierarchy) {
93   for (int icontour=0; icontour >= 0; icontour = hierarchy[icontour][0]) {
94     fishblobb newblobb;
95
96     if ( contours[icontour].size() < 5 )
97       continue;
98
99     newblobb.circum = contours[icontour].size();
100     if ( newblobb.circum < normalPrefs.contours_mincircum || newblobb.circum > normalPrefs.contours_maxcircum )
101       continue;
102
103     newblobb.area = contourArea(contours[icontour]);
104     if ( newblobb.area < normalPrefs.contours_minarea || newblobb.area > normalPrefs.contours_maxarea )
105       continue;
106
107     RotatedRect minEllipse = fitEllipse(contours[icontour]);
108     newblobb.longaxis = minEllipse.size.height;
109     newblobb.shortaxis = minEllipse.size.width;
110     if ( newblobb.shortaxis < normalPrefs.contours_minshortaxis || newblobb.shortaxis > normalPrefs.contours_maxshortaxis
111          || newblobb.longaxis < normalPrefs.contours_minlongaxis || newblobb.longaxis > normalPrefs.contours_maxlongaxis )
112       continue;
113      
114     newblobb.x = minEllipse.center.x;
115     newblobb.y = minEllipse.center.y;
116      
117     RotatedRect minRect;
118     Mat mFish = Mat::zeros(Props.height, Props.width, CV_8U);
119     Moments fishMom;
120     Point2f massCenter;
121     Point2f recHead, recTail;
122         
123     newblobb.identified = false;
124     newblobb.points = &contours[icontour];
125
126     newblobb.contourid = icontour;
127
128     /* find head/tail sensitive angle
129        maybe the method could be improved */
130     drawContours( mFish, contours, icontour, Scalar(255, 255, 255), CV_FILLED );
131     fishMom = moments(mFish);
132       
133     massCenter = Point2f( fishMom.m10/fishMom.m00 , fishMom.m01/fishMom.m00 );
134     minRect = minAreaRect(contours[ icontour ]);
135       
136     recHead = Point2f( minRect.center.x + newblobb.longaxis*cos(minEllipse.angle * M_PI / 180 + M_PI/2), 
137                        minRect.center.y + newblobb.longaxis*sin(minEllipse.angle * M_PI / 180 + M_PI/2) );
138     recTail = Point2f( minRect.center.x + newblobb.longaxis*cos(minEllipse.angle * M_PI / 180 - M_PI/2), 
139                        minRect.center.y + newblobb.longaxis*sin(minEllipse.angle * M_PI / 180 - M_PI/2) );
140       
141     if ( pow( recHead.x - massCenter.x, 2 ) + pow( recHead.y - massCenter.y, 2 ) > pow( recTail.x - massCenter.x, 2 ) + pow( recTail.y - massCenter.y, 2 ) ) {
142       newblobb.angle = minEllipse.angle + 180;
143         
144       if (newblobb.angle > 360)
145         newblobb.angle = newblobb.angle - 360;
146     }
147     else {
148       newblobb.angle = minEllipse.angle;
149     }
150
151     newblobb.tail = Point2f( minEllipse.center.x + newblobb.longaxis*cos(newblobb.angle * M_PI / 180 - M_PI/2),
152                              minEllipse.center.y + newblobb.longaxis*sin(newblobb.angle * M_PI / 180 - M_PI/2) );
153     newblobb.head = Point2f( minEllipse.center.x + newblobb.longaxis*cos(newblobb.angle * M_PI / 180 + M_PI/2), 
154                              minEllipse.center.y + newblobb.longaxis*sin(newblobb.angle * M_PI / 180 + M_PI/2) );
155
156     fishblobbs.push_back(newblobb);
157   }
158 }
159
160 void tagCare(vector<tag>& tags) {
161   while ( tags.size() > Prefs.manyfish )
162     tags.pop_back();
163
164   while ( tags.size() < Prefs.manyfish ) {
165     tag tmp;
166     tmp.x = 0;
167     tmp.y = 0;
168     tmp.propZ = 0;
169     tmp.lastseen = frametime;
170
171     switch ( tags.size() ) {
172     case 0:
173       tmp.color = Scalar( 0, 255, 255 );
174       break;
175     case 1:
176       tmp.color = Scalar( 255, 0, 255 );
177       break;
178     case 2:
179       tmp.color = Scalar( 255, 255, 0 );
180       break;
181     case 3:
182       tmp.color = Scalar( 255, 0, 0 );
183       break;
184     case 4:
185       tmp.color = Scalar( 0, 255, 0 );
186       break;
187     case 5:
188       tmp.color = Scalar( 0, 0, 255 );
189       break;
190     default:
191       tmp.color = Scalar( rand()%255, rand()%255, rand()%255 );
192     }
193
194     tmp.angle = 0;
195     tmp.area = 0;
196     tmp.circum = 0;
197     tmp.shortaxis = 0;
198     tmp.longaxis = 0;
199     tmp.virgin = true;
200     tags.push_back(tmp);
201   }
202 }
203
204 void matchBlobbsNTags(vector<tag>& tags, vector<fishblobb>& fishblobbs) {
205   vector<blobbcost> blobbcosts;
206
207   for (int iblobb = 0; iblobb < fishblobbs.size(); iblobb++) {
208     double sepcost[tags.size()][6];
209     double maxcost[6];
210     blobbcost blobbCost;
211
212     blobbCost.blobbid = iblobb;
213
214     for (int imax = 0; imax < 6; imax++)
215       maxcost[imax] = 0;
216
217     for (int itag = 0; itag < tags.size(); itag++) {
218       double anglediff = fabs (fishblobbs[iblobb].angle - tags[itag].angle);
219             
220       if ( anglediff > 180 )
221         anglediff = 360 - anglediff;
222
223       sepcost[itag][SEPCOST_ANGLE] = anglediff;
224
225       sepcost[itag][SEPCOST_AREA] = fabs ( fishblobbs[iblobb].area - tags[itag].area );
226       sepcost[itag][SEPCOST_CIRCUM] = fabs ( fishblobbs[iblobb].circum - tags[itag].circum );
227       sepcost[itag][SEPCOST_SHORTAXIS] = fabs ( fishblobbs[iblobb].shortaxis - tags[itag].shortaxis );
228       sepcost[itag][SEPCOST_LONGAXIS] = fabs ( fishblobbs[iblobb].longaxis - tags[itag].longaxis );
229       sepcost[itag][SEPCOST_DIST] = sqrt( pow(tags[itag].x - fishblobbs[iblobb].x, 2) + pow(tags[itag].y - fishblobbs[iblobb].y, 2) );
230
231       for (int imax = 0; imax < 6; imax++)
232         if ( maxcost[imax] < sepcost[itag][imax] )
233           maxcost[imax] = sepcost[itag][imax];
234     }
235
236     for (int itag = 0; itag < tags.size(); itag++) {
237       tagcost newTagCost;
238             
239       newTagCost.tagid = itag;
240       newTagCost.tagcost = 0;
241
242       if ( sepcost[itag][SEPCOST_DIST] > ( frametime - tags[itag].lastseen )*normalPrefs.contours_maxspeed 
243            || sepcost[itag][SEPCOST_ANGLE] > ( frametime - tags[itag].lastseen )*normalPrefs.contours_maxrot ) 
244         continue;
245
246       for (int isep = 0; isep < 6; isep++)
247         newTagCost.tagcost += sepcost[itag][isep] / maxcost[isep];
248
249       blobbCost.tagcosts.push_back( newTagCost );                                               
250     }
251
252     if ( blobbCost.tagcosts.size() > 0 ) { 
253       sort( blobbCost.tagcosts.begin(), blobbCost.tagcosts.end(), sort_tags );
254       blobbcosts.push_back( blobbCost );
255     }
256   }
257         
258   for (int itag = 0; itag < tags.size() && blobbcosts.size() > 0; itag++) {
259     sort( blobbcosts.begin(), blobbcosts.end(), sort_blobbs );
260
261     int tagnow = blobbcosts[0].tagcosts[0].tagid;
262
263     tags[tagnow].virgin = false;
264     tags[tagnow].x = fishblobbs[ blobbcosts[0].blobbid ].x;
265     tags[tagnow].y = fishblobbs[ blobbcosts[0].blobbid ].y;
266     tags[tagnow].shortaxis = fishblobbs[ blobbcosts[0].blobbid ].shortaxis;
267     tags[tagnow].longaxis =  fishblobbs[ blobbcosts[0].blobbid ].longaxis;
268     tags[tagnow].angle = fishblobbs[ blobbcosts[0].blobbid ].angle;
269     tags[tagnow].lastseen = frametime;
270     tags[tagnow].area = fishblobbs[ blobbcosts[0].blobbid ].area;
271     tags[tagnow].circum = fishblobbs[ blobbcosts[0].blobbid ].circum;
272     tags[tagnow].head = fishblobbs[ blobbcosts[0].blobbid ].head;
273     tags[tagnow].tail = fishblobbs[ blobbcosts[0].blobbid ].tail;
274
275     blobbcosts.erase( blobbcosts.begin() );
276
277     for (int iblobb = 0; iblobb < blobbcosts.size(); iblobb++) {
278             
279       for (int itag = 0; itag < blobbcosts[iblobb].tagcosts.size(); itag++)
280         if ( blobbcosts[iblobb].tagcosts[itag].tagid == tagnow ) {
281           blobbcosts[iblobb].tagcosts.erase( blobbcosts[iblobb].tagcosts.begin() + itag );
282           break;
283         }
284            
285       if ( blobbcosts[iblobb].tagcosts.size() == 0 ) {
286         blobbcosts.erase( blobbcosts.begin() + iblobb );
287         iblobb--;
288       }
289     }
290   }
291 }
292
293 void drawTags (const vector<tag>& tags) {
294   RotatedRect rRect;
295   Point2f vertices[4];
296
297   for (int itag = 0; itag < tags.size(); itag++) {
298     double recside = sqrt( kitvalue( normalPrefs.contours_minarea, normalPrefs.contours_maxarea) );
299     rRect = RotatedRect( Point(tags[itag].x, tags[itag].y), Size2f(recside, recside), tags[itag].angle );
300     rRect.points(vertices);
301     for (int i = 0; i < 4; i++)
302       line(mContours, vertices[i], vertices[(i+1)%4], Scalar(255,255,255), Props.diagonal/1000.0);
303
304     ellipse(mContours, Point(tags[itag].x, tags[itag].y), 
305             Size( kitvalue( normalPrefs.contours_minshortaxis, normalPrefs.contours_maxshortaxis ), 
306                   kitvalue( normalPrefs.contours_minlongaxis, normalPrefs.contours_maxlongaxis ) ), 
307             tags[itag].angle, 0, 360, Scalar(255,255,255), Props.diagonal/1000.0, 8);
308
309     circle(mContours, Point(tags[itag].x, tags[itag].y), kitvalue( normalPrefs.contours_mincircum, normalPrefs.contours_maxcircum )/(2*M_PI), Scalar(255,255,255), 
310            Props.diagonal/1000.0, 8);
311       
312     if ( ! tags[itag].virgin ) {
313       float searchradius = ( frametime - tags[itag].lastseen )*normalPrefs.contours_maxspeed;
314       float searchangle = ( frametime - tags[itag].lastseen )*normalPrefs.contours_maxrot;
315       ellipse(mContours, Point(tags[itag].x, tags[itag].y), Size(searchradius, searchradius), tags[itag].angle, 90-searchangle, 90+searchangle, 
316               tags[itag].color, Props.diagonal/1000.0, 8);
317       
318       rRect = RotatedRect( Point(tags[itag].x, tags[itag].y), Size2f(sqrt(tags[itag].area), sqrt(tags[itag].area)), tags[itag].angle );
319       rRect.points(vertices);
320       for (int i = 0; i < 4; i++)
321         line(mContours, vertices[i], vertices[(i+1)%4], tags[itag].color, 2);
322       
323       ellipse(mContours, Point(tags[itag].x, tags[itag].y), 
324               Size( tags[itag].shortaxis, tags[itag].longaxis ), tags[itag].angle, 0, 360, tags[itag].color, 2*Props.diagonal/1000.0, 8);
325
326       circle(mContours, Point(tags[itag].x, tags[itag].y), tags[itag].circum / (2*M_PI), tags[itag].color, 2*Props.diagonal/1000.0, 8);
327       
328       circle(mContours, tags[itag].head, 5*Props.diagonal/1000.0, tags[itag].color, -1, 8);
329     }
330   }
331 }
332
333 void tracking_locateTags (vector<tag>& tags, Mat combinedmask_contour) {
334   mContours = Mat::zeros(Props.height, Props.width, CV_8UC3);
335
336   if (enablecontours) {
337     vector< vector<Point> > contours;
338     vector<Vec4i> hierarchy;
339     vector <fishblobb> fishblobbs;
340
341     findContours( combinedmask_contour, contours, hierarchy, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE );
342
343     if ( contours.size() > 0 ) {
344
345       for(int idx = 0 ; idx >= 0; idx = hierarchy[idx][0] )
346         drawContours(mContours, contours, idx, Scalar(255, 255, 255), CV_FILLED, 8, hierarchy);
347
348       findFishBlobbs(fishblobbs, contours, hierarchy);
349
350       for ( int iblobb = 0; iblobb < fishblobbs.size(); iblobb++) {
351         Scalar color( rand()%255, rand()%255, rand()%255 );
352         drawContours( mContours, contours, fishblobbs[iblobb].contourid, color, CV_FILLED, 8, hierarchy );
353       }
354
355       tagCare(tags);
356       matchBlobbsNTags(tags, fishblobbs);
357
358     }
359
360     drawTags(tags);
361   }
362 }