stickletrack.cpp

   1 #include "opencv2/highgui/highgui.hpp"
   2 #include "opencv2/imgproc/imgproc.hpp"
   3 #include <iostream>
   4 #include <algorithm>
   5 #include <functional>
   6 #include <sys/stat.h>
   7 #include <fstream>
   8 #include <sstream>
   9 #include <iterator>
  10 #include <ctime>
  11 #include <iomanip>
  12
  13 // needed for _mkdir
  14 #if defined(_WIN32)
  15 #include <direct.h>
  16 #endif
  17
  18 #include "masking.h"
  19 #include "stickletrack.h"
  20 #include "tracking.h"
  21
  22 #define BACKSECONDS 30
  23
  24 double frametime = 0;
  25 int framenum = 0;
  26
  27 prefs Prefs;
  28 props Props;
  29
  30 normalprefs normalPrefs;
  31
  32 static VideoCapture capture;
  33
  34 static bool stopvideo=false;
  35 static bool leftbuttondown=false;
  36 static Point2f startClickPos;
  37 static int ilatetagintime=-1, iearlytagintime=-1;
  38 static int startClickFrame;
  39 static bool stoppedbefore;
  40 static bool dooneframe=false;
  41 static vector<tag> tags;
  42 static int tagsselected=0;
  43 static int nearestTags[2];
  44 static vector<tag> **tagintime;
  45 static int tagintimeidx=0;
  46 static int moresleep=0;
  47 static int gotoframe=0;
  48 static int maxoutputfps;
  49 static double rescalingfactor;
  50 static ofstream tanjaLogFile;
  51
  52 void genBaseDir() {
  53 #if defined(_WIN32)
  54   Props.basedir = getenv("HOMEDRIVE");
  55   Props.basedir += getenv("HOMEPATH");
  56   Props.basedir += "\.stickletrack";
  57 #else
  58   Props.basedir = getenv("HOME");
  59   Props.basedir += "/.stickletrack";
  60 #endif
  61
  62   cout << "Using " << Props.basedir << " for data output." << endl;
  63 }
  64
  65 void beforeExit() {
  66   cout << "Exitting ..." << endl;
  67
  68 #if defined(_WIN32)
  69   _mkdir(Props.basedir.c_str());
  70   _mkdir((Props.basedir+"/prefs").c_str());
  71 #else
  72   mkdir( Props.basedir.c_str(), 0777 );
  73   mkdir( (Props.basedir+"/prefs").c_str(), 0777 );
  74 #endif
  75
  76   cout << "Saving preferences ..." << endl;
  77   int scissorpointsize = Prefs.scissorpoints->size();
  78   ofstream prefFile( (Props.basedir+"/prefs/"+Props.videohash+".prf").c_str(), ios::binary );
  79   prefFile.write((char*)&Prefs, sizeof(Prefs));
  80   prefFile.write((char*)&scissorpointsize, sizeof(int));
  81
  82   for (int ipoint = 0; ipoint < scissorpointsize; ipoint++)
  83     prefFile.write( (char*)&(*Prefs.scissorpoints)[ipoint], sizeof(Point) );
  84
  85   prefFile.close();
  86 }
  87
  88 void computeNormalizedPrefs() {
  89   normalPrefs.contours_minshortaxis = Prefs.contours_minshortaxis * Props.diagonal / 1000.0;
  90   normalPrefs.contours_maxshortaxis = Prefs.contours_maxshortaxis * Props.diagonal / 1000.0;
  91   normalPrefs.contours_minlongaxis = Prefs.contours_minlongaxis * Props.diagonal / 500.0;
  92   normalPrefs.contours_maxlongaxis = Prefs.contours_maxlongaxis * Props.diagonal / 500.0;
  93
  94   normalPrefs.contours_minarea = pow(Prefs.contours_minarea,2) * Props.width * Props.height / 1000000.0;
  95   normalPrefs.contours_maxarea = pow(Prefs.contours_maxarea,2) * Props.width * Props.height / 1000000.0;
  96
  97   normalPrefs.contours_mincircum = Prefs.contours_mincircum * Props.diagonal / 500.0;
  98   normalPrefs.contours_maxcircum = Prefs.contours_maxcircum * Props.diagonal / 500.0;
  99
 100   normalPrefs.contours_maxspeed = Prefs.contours_maxspeed * Props.diagonal / 100.0;
 101
 102   normalPrefs.contours_maxrot = 10 * Prefs.contours_maxrot;
 103
 104   normalPrefs.halfdecay = 10.0 * Prefs.halfdecay / Props.fps;
 105 }
 106
 107 void trackbarCallbackUpdateNormPrefs (int trackpos, void *userdata) {
 108   computeNormalizedPrefs();
 109 }
 110
 111 void drawTimes(Mat& mContours) {
 112   stringstream alles;
 113
 114   alles.precision(2);
 115
 116   alles << framenum << " : " << fixed << (framenum)/(double)Props.fps << "s : +" << (double)moresleep/Props.fps << "ms";
 117
 118   string text = alles.str();
 119   int fontFace = FONT_HERSHEY_SCRIPT_SIMPLEX;
 120
 121   double fontScale = Props.diagonal/1000.0;
 122   int thickness = 3;
 123
 124   int baseline=0;
 125   Size textSize = getTextSize(text, fontFace,
 126                               fontScale, thickness, &baseline);
 127
 128   Point textOrg(10, 10 + textSize.height);
 129
 130   rectangle(mContours, Point(0, 0),
 131             Point(Props.width, 20+textSize.height),
 132             Scalar(0,0,255), -1);
 133
 134   putText(mContours, text, textOrg, fontFace, fontScale,
 135           Scalar::all(255), thickness, 8);
 136 }
 137
 138 void readCreatePrefs() {
 139   ifstream prefFile( (Props.basedir+"/prefs/"+Props.videohash+".prf").c_str(), ios::binary );
 140   if ( prefFile.is_open() ) {
 141     int scissorpointsize;
 142
 143     prefFile.read((char*)&Prefs, sizeof(Prefs));
 144     prefFile.read((char*)&scissorpointsize, sizeof(int));
 145
 146     Prefs.scissorpoints = new vector<Point>;
 147     for (int ipoint = 0; ipoint < scissorpointsize; ipoint++) {
 148       Point tmpPoint;
 149       prefFile.read((char*)&tmpPoint, sizeof(Point));
 150       Prefs.scissorpoints->push_back(tmpPoint);
 151     }
 152
 153     prefFile.close();
 154
 155     cout << "Loaded saved preferences." << endl;
 156   }
 157   else {
 158     Prefs.scissorpoints = new vector<Point>;
 159   }
 160 }
 161
 162 void loadDefaultPrefs () {
 163   Prefs.halfdecay = 100;
 164   Prefs.forethreshold = 0;
 165   Prefs.mincolor[0] = 0; Prefs.mincolor[1] = 0; Prefs.mincolor[2] = 0;
 166   Prefs.maxcolor[0] = 255; Prefs.maxcolor[1] = 255; Prefs.maxcolor[2] = 255;
 167   Prefs.manyfish = 3;
 168   Prefs.contours_minshortaxis = 0;
 169   Prefs.contours_maxshortaxis = 100;
 170   Prefs.contours_minlongaxis = 0;
 171   Prefs.contours_maxlongaxis = 100;
 172   Prefs.contours_minarea = 0;
 173   Prefs.contours_maxarea = 100;
 174   Prefs.contours_mincircum = 0;
 175   Prefs.contours_maxcircum = 100;
 176   Prefs.contours_maxspeed = 100;
 177   Prefs.contours_maxrot = 100;
 178 }
 179
 180 unsigned long frameHash(const Mat& frame) {
 181   unsigned long hash = 5381;
 182   int c;
 183   unsigned char *str = (unsigned char*)frame.data;
 184
 185   while (c = *str++) {
 186     hash = ((hash << 5) + hash) + c;
 187   }
 188
 189   return hash;
 190 }
 191
 192 int timeWarp(int x, int y) {
 193   int horiMove = startClickPos.x - x;
 194   int wannago = startClickFrame + 10*horiMove*Props.fps/Props.width;
 195
 196   int maxgo = min(ilatetagintime, Props.totframes);
 197   int mingo = max(1, ilatetagintime - BACKSECONDS*Props.fps + 2);
 198
 199   if (wannago > maxgo)
 200     return maxgo;
 201   else if (wannago < mingo)
 202     return mingo;
 203   else
 204     return wannago;
 205 }
 206
 207 int rotatingTime( int whereyouwant ) {
 208   if ( whereyouwant < 0 )
 209     return BACKSECONDS*Props.fps + whereyouwant;
 210   else
 211     return whereyouwant;
 212 }
 213
 214 void mouseTracking(int evt, int x, int y, int flags, void* param){
 215   if ( stopvideo && evt == CV_EVENT_LBUTTONDBLCLK ) {
 216     tagsselected++;
 217
 218     double mindist = -1;
 219     for (int itag = 0; itag < tags.size(); itag++) {
 220       double dist = pow(tags[itag].x - x, 2) + pow(tags[itag].y - y, 2);
 221       if ( dist < mindist || mindist == -1 ) {
 222         nearestTags[tagsselected-1] = itag;
 223         mindist = dist;
 224       }
 225     }
 226
 227     if (tagsselected == 2) {
 228       if ( nearestTags[0] < tags.size() && nearestTags[0] >= 0
 229            && nearestTags[1] < tags.size() && nearestTags[1] >= 0
 230            && nearestTags[0] != nearestTags[1] ) {
 231
 232         tag tmp;
 233
 234         tmp = tags[nearestTags[0]];
 235         tags[nearestTags[0]] = tags[nearestTags[1]];
 236         tags[nearestTags[1]] = tmp;
 237         tags[nearestTags[1]].color = tags[nearestTags[0]].color;
 238         tags[nearestTags[0]].color = tmp.color;
 239
 240         for (int itagintime = framenum-1; itagintime <= ilatetagintime; itagintime++) {
 241           tmp = (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[0]];
 242
 243           (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[0]] =
 244             (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[1]];
 245
 246           (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[1]] = tmp;
 247
 248           (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[1]].color =
 249             (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[0]].color;
 250
 251           (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-itagintime) - 1) ])[nearestTags[0]].color = tmp.color;
 252         }
 253       }
 254       tagsselected = 0;
 255
 256       gotoframe = framenum;
 257       dooneframe = true;
 258     }
 259   }
 260
 261   if ( evt == CV_EVENT_RBUTTONUP ) {
 262     stopvideo = (stopvideo+1)%2;
 263   }
 264
 265   if (evt==CV_EVENT_LBUTTONDOWN) {
 266     leftbuttondown = true;
 267     startClickPos = Point2f(x,y);
 268     startClickFrame = framenum;
 269
 270     if (stopvideo)
 271       stoppedbefore = true;
 272     else
 273       stoppedbefore = false;
 274
 275     stopvideo = true;
 276   }
 277
 278   if (evt==CV_EVENT_LBUTTONUP) {
 279     leftbuttondown = false;
 280
 281     int vertiMove = startClickPos.y - y;
 282
 283     if ( ! stoppedbefore )
 284       stopvideo = false;
 285  }
 286
 287   if ( leftbuttondown && startClickPos.x - x != 0  && ! (startClickPos.x > Props.width*5.0/6.0 && x > Props.width*5.0/6.0) ) {
 288     gotoframe = timeWarp(x, y);
 289     dooneframe = true;
 290   }
 291
 292   if ( leftbuttondown && startClickPos.y - y != 0) {
 293     int vertiMove = startClickPos.y - y;
 294
 295     if ( startClickPos.x > Props.width*5.0/6.0 && x > Props.width*5.0/6.0 ) {
 296       moresleep -= vertiMove*10000/Props.height;
 297       if ( moresleep < 0 )
 298         moresleep = 0;
 299
 300       drawTimes(tracking_getFrame());
 301     }
 302   }
 303 }
 304
 305 void writeTanjaLog (int writeframenum, vector<tag> *tagstowrite) {
 306   stringstream outline;
 307
 308   if ( writeframenum%(Props.fps/maxoutputfps) == 0 ) {
 309     if ( tagstowrite->size() >= Prefs.manyfish ) {
 310       outline << (double)writeframenum/Props.fps;
 311       for ( int itag = 0; itag < Prefs.manyfish; itag++) {
 312         outline << "," << (*tagstowrite)[itag].x << "," << (*tagstowrite)[itag].y << "," << (double)writeframenum/Props.fps-(*tagstowrite)[itag].lastseen;
 313       }
 314       outline << endl;
 315       tanjaLogFile.write( outline.str().c_str(), outline.str().size() );
 316       tanjaLogFile.flush();
 317     }
 318   }
 319 }
 320
 321 void openTanjaLog () {
 322   stringstream tanjalogfilename;
 323
 324 #if defined(_WIN32)
 325   _mkdir(Props.basedir.c_str());
 326   _mkdir((Props.basedir+"/logs").c_str());
 327 #else
 328   mkdir( Props.basedir.c_str(), 0777 );
 329   mkdir( (Props.basedir+"/logs").c_str(), 0777 );
 330 #endif
 331
 332 #if defined(_WIN32)
 333   _mkdir(Props.basedir.c_str());
 334   _mkdir((Props.basedir+"/logs/"+Props.videohash).c_str());
 335 #else
 336   mkdir( Props.basedir.c_str(), 0777 );
 337   mkdir( (Props.basedir+"/logs/"+Props.videohash).c_str(), 0777 );
 338 #endif
 339
 340   tanjalogfilename << Props.basedir << "/logs/"  << Props.videohash << "/" << time(0) << ".dat";
 341   tanjaLogFile.open( tanjalogfilename.str().c_str() );
 342
 343   if ( ! tanjaLogFile.is_open() ) {
 344     cerr << "Could not open tanjalogfile " << tanjalogfilename.str() << "!" << endl;
 345     exit(1);
 346   }
 347
 348   tanjaLogFile.write("# STICKLETRACK - TANJALOG\n", 26);
 349   tanjaLogFile.write("# format: <time in s>,<tag_1 xpos>,<tag_1 ypos>,<tag_1 losttime>,...,<tag_n xpos>,<tag_n ypos>,<tag_n losttime>\n", 112);
 350
 351   tanjaLogFile.flush();
 352 }
 353
 354 int process(VideoCapture& capture) {
 355   int n = 0;
 356   char filename[200];
 357   bool pleaseExit = false;
 358
 359   namedWindow("original", CV_WINDOW_KEEPRATIO);
 360
 361   Mat frame, origframe, combinedmask;
 362
 363   Props.fps = capture.get(CV_CAP_PROP_FPS);
 364   Props.width = capture.get(CV_CAP_PROP_FRAME_WIDTH)*rescalingfactor;
 365   Props.height = capture.get(CV_CAP_PROP_FRAME_HEIGHT)*rescalingfactor;
 366   Props.totframes = capture.get(CV_CAP_PROP_FRAME_COUNT);
 367   Props.diagonal = sqrt( pow(Props.width, 2) + pow(Props.height, 2) );
 368
 369   Mat frameintime[BACKSECONDS*Props.fps];
 370
 371   tagintime = new vector<tag>*[BACKSECONDS*Props.fps];
 372
 373   for (int iback = 0; iback < BACKSECONDS*Props.fps; iback++)
 374     tagintime[iback] = new vector<tag>;
 375
 376   capture >> frame;
 377
 378   if (frame.empty())
 379     exit(1);
 380
 381   stringstream tmphash;
 382   tmphash << frameHash(frame);
 383   Props.videohash = tmphash.str();
 384
 385   loadDefaultPrefs();
 386   readCreatePrefs();
 387
 388   computeNormalizedPrefs();
 389
 390   openTanjaLog();
 391
 392   masking_init();
 393   tracking_init();
 394
 395   cvSetMouseCallback("contours_picture", mouseTracking, 0);
 396
 397   capture.set(CV_CAP_PROP_POS_FRAMES, 0);
 398
 399   for (; !pleaseExit;) {
 400     if ( !stopvideo || dooneframe ) {
 401       framenum = gotoframe;
 402
 403       dooneframe = false;
 404
 405       frametime = (framenum-1.0) / Props.fps;
 406
 407       if ( framenum <= ilatetagintime ) {
 408         origframe = frameintime[ rotatingTime(tagintimeidx - (ilatetagintime-framenum) - 2) ];
 409         masking_getCombinedMask(origframe, combinedmask);
 410
 411         tags.clear();
 412         for (int itag = 0; itag < tagintime[ rotatingTime( tagintimeidx - 2 - (ilatetagintime-framenum) ) ]->size(); itag++)
 413           tags.push_back( (*tagintime[ rotatingTime(tagintimeidx - (ilatetagintime-framenum) - 2) ])[itag] );
 414
 415         tracking_locateTags (tags, combinedmask);
 416       }
 417       else {
 418         capture >> frame;
 419
 420         if (frame.empty())
 421           break;
 422
 423         frame.convertTo(origframe, CV_32FC3);
 424         resize(origframe, origframe, Size(0,0), rescalingfactor, rescalingfactor);
 425
 426         masking_getCombinedMask(origframe, combinedmask);
 427
 428         tracking_locateTags (tags, combinedmask);
 429
 430         tagintime[tagintimeidx]->clear();
 431         for (int itag = 0; itag < tags.size(); itag++)
 432           tagintime[tagintimeidx]->push_back(tags[itag]);
 433
 434         frameintime[tagintimeidx] = origframe;
 435
 436         ilatetagintime = framenum;
 437         tagintimeidx = (tagintimeidx+1)%(BACKSECONDS*Props.fps);
 438
 439         if ( ilatetagintime >= BACKSECONDS*Props.fps-1 )
 440           writeTanjaLog( framenum - (BACKSECONDS*Props.fps-1), tagintime[ tagintimeidx ] );
 441       }
 442
 443       gotoframe = framenum + 1;
 444
 445       imshow("original", origframe/255.0);
 446     }
 447
 448     drawTimes(tracking_getFrame());
 449
 450     if ( tagsselected == 1 ) {
 451       circle( tracking_getFrame(), Point2f(tags[nearestTags[0]].x, tags[nearestTags[0]].y), Props.diagonal / 100.0, Scalar(0,0,255), -1, 8 );
 452     }
 453
 454     tracking_showFrame();
 455
 456     char key;
 457
 458     if (!stopvideo)
 459       key = (char)waitKey(5 + (double)moresleep/Props.fps);
 460     else
 461       key = (char)waitKey(5);
 462
 463     switch (key) {
 464     case 'q':
 465       pleaseExit = true;
 466       break;
 467     default:
 468       break;
 469     }
 470
 471   }
 472
 473   if ( ilatetagintime >= BACKSECONDS*Props.fps-1 ) {
 474     int timeidx = (tagintimeidx+1)%(BACKSECONDS*Props.fps);
 475     for ( int itime = 0; itime < BACKSECONDS*Props.fps-1; itime++ ) {
 476       writeTanjaLog( framenum+1+itime - (BACKSECONDS*Props.fps-1), tagintime[ timeidx ] );
 477       timeidx = (timeidx+1)%(BACKSECONDS*Props.fps);
 478     }
 479   }
 480   else {
 481     for ( int itime = 0; itime < ilatetagintime+1; itime++ )
 482       writeTanjaLog( itime, tagintime[ itime ] );
 483   }
 484
 485   tanjaLogFile.close();
 486
 487   beforeExit();
 488
 489   cout << "Bye bye." << endl;
 490
 491   return 0;
 492 }
 493
 494 int main(int ac, char** av) {
 495   if ( ac != 4 ) {
 496     cout << "Usage: stickletrack <maxoutputfps> <rescalingfactor> <videofilename>" << endl;
 497     exit(0);
 498   }
 499
 500   maxoutputfps = atoi( av[1] );
 501   rescalingfactor = atof( av[2] );
 502
 503   capture.open( av[3] );
 504
 505   genBaseDir();
 506
 507    if (!capture.isOpened()) {
 508     cerr << "Failed to open a video device or video file!\n" << endl;
 509     return 1;
 510   }
 511
 512    cout << "Exit with q if you want to save your data!!!" << endl;
 513
 514   return process(capture);
 515 }
 516