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