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