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