preTrackingAnnotation.py

#import matplotlib.pyplot as plt
#import matplotlib.widgets as widgets
import numpy as np
import cv2
import copy
import sys
if sys.version_info >= (3, 0):
   import tkinter
   # import filedialog
   from tkinter import filedialog as tkFileDialog
else:
   import Tkinter # python 3: tkinter
   import tkFileDialog # python 3: filedialog

#import trackFlies as tf
import os.path
import glob
import skvideo.io

import scipy.io as sio
import argparse

def updateImage(x, y):
   global arenaX, arenaY, rectRadius, rectCenterX, rectCenterY, showFrame
   
   xdelta = abs(arenaX-x);
   ydelta = abs(arenaY-y);
   delta = max(xdelta,ydelta);

   rectRadius = delta/2;
   rectCenterX = arenaX;
   rectCenterY = arenaY;
   
   img = copy.copy(showFrame);
   cv2.circle(img, (arenaX, arenaY), delta, (0, 255, 0), 3)
   cv2.imshow('display',img)

def areaCallback(event, x, y, flags, param):
   global currentlyDrawing, arenaX, arenaY

   if event == cv2.EVENT_LBUTTONDOWN and currentlyDrawing == False:
      arenaX = x;
      arenaY = y;

   if event == cv2.EVENT_LBUTTONDOWN:
      currentlyDrawing = True;
   elif event == cv2.EVENT_LBUTTONUP:
      currentlyDrawing = False;
   
   if currentlyDrawing:
      updateImage(x,y);

def updateCenterImage(x, y):
   global arenaX, arenaY, showFrame
   
   img = copy.copy(showFrame);
   cv2.circle(img, (x, y), 5, (0, 255, 0), 3)
   cv2.imshow('display',img)

def centerCallback(event, x, y, flags, param):
   global currentlyDrawing, arenaX, arenaY

   if event == cv2.EVENT_LBUTTONDOWN and currentlyDrawing == False:
      arenaX = x;
      arenaY = y;

   if event == cv2.EVENT_LBUTTONDOWN:
      currentlyDrawing = True;
      arenaX = x;
      arenaY = y;
   if event == cv2.EVENT_LBUTTONUP:
      currentlyDrawing = False;

   if currentlyDrawing:
      updateCenterImage(x,y);

def trackerCallback(position):
   global angle
   
   M = cv2.getRotationMatrix2D((width/4,height/4),position-180,1);
   angle = position;
   img = copy.copy(showFrame);
   img = cv2.warpAffine(img,M,(int(width/2),int(height/2)));
   cv2.imshow('display',img);

def findMovieFile(fName):
   print(fName)
   if fName[len(fName)-4:len(fName)] == '.mp4' or fName[len(fName)-4:len(fName)] == '.avi':
      return fName
   else:
      nName = glob.glob(fName + '/*.avi')
      if len(nName) > 0:
         return nName[0]
      else:
         nName = glob.glob(fName + '/*.mp4')
         # avoid collision with the *Cmp.mp4 files generated by Pip's tracker - for some reason they will be the first match...
         # this cleans the list of MP4 files of the ones containing "Cmp.mp4"
         nName = [thisName for thisName in nName if "Cmp.mp4" not in thisName]
         if len(nName) > 0:
            return nName[0]
         else:
            return None


# these are the global variables that we want to assign
currentlyDrawing = False;
arenaX = 0
arenaY = 0

delta = 0
rectCenterX = 0
rectCenterY = 0
rectRadius = 0
angle = 0
micCenters = []

centerX = 0
centerY = 0
radius = 0

newAnnotation = True
applyAll = False

# step zero: choose a directory that you wish to annotate all of
parser = argparse.ArgumentParser()
parser.add_argument('initialdir',type=str,default='.', help='Movie to analyze')
pargs = parser.parse_args()

dirName = tkFileDialog.askdirectory(initialdir=pargs.initialdir)
names = sorted(glob.glob(dirName + '/*'))

cv2.namedWindow('display')

for movieName in names:
   if not os.path.isdir(movieName):
      continue

   # step one: load fly video
   print(movieName)
   fileName = findMovieFile(movieName)
   print(fileName)
   if fileName is None:
      continue

   if not applyAll:


      try:
         videoInfo = skvideo.io.ffprobe(fileName)
         width = int(videoInfo['video']['@width'])
         height = int(videoInfo['video']['@height'])

         videodata = skvideo.io.vreader(fileName)

         frame = next(videodata)
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
      except:
         continue

      cv2.resizeWindow('display',width, height)

      frame2 = copy.copy(frame);
      cv2.putText(frame2, 'Should we use this annotation? (y/n/a/q)', (50,50), cv2.FONT_HERSHEY_PLAIN, 2, 255)
      showFrame = cv2.resize(frame2,(int(width/2),int(height/2)));
      cv2.circle(showFrame,(centerX,centerY),int(radius),(0, 255, 0), 3)
      cv2.imshow('display',showFrame)
      key = cv2.waitKey(0)
      if key == ord('y'):
         newAnnotation = False
      elif key == ord('n'):
         newAnnotation = True
         arenaX = 0
         arenaY = 0

         delta = 0
         rectCenterX = 0
         rectCenterY = 0
         rectRadius = 0
         angle = 0
         micCenters = []

         centerX = 0
         centerY = 0
         radius = 0
      elif key == ord('a'):
         applyAll = True
         newAnnotation = False
      elif key == ord('q'):
         exit()

      print('got key ' + chr(key))


      if newAnnotation:
         # for each of these, let the user press 'x' to reset, 'q' to quit at any time

         # step two: choose area that contains fly arena
         resetThis = True
         while resetThis == True:
            resetThis = False

            print((width,height))
            frame2 = copy.copy(frame);
            cv2.putText(frame2, 'Choose arena (press enter when done)', (50,50), cv2.FONT_HERSHEY_PLAIN, 2, 255)
            showFrame = cv2.resize(frame2,(int(width/2),int(height/2)));
            cv2.imshow('display',showFrame)
            cv2.setMouseCallback('display',centerCallback)
            key = cv2.waitKey(0)
            if key == 'x':
               resetThis = True
               break
            if key == 'q':
               exit()

            centerX, centerY = arenaX, arenaY;

            # step three: select arena width
            frame2 = copy.copy(frame);
            cv2.putText(frame2, 'Select LEFT boundary', (50,50), cv2.FONT_HERSHEY_PLAIN, 2, 255)
            showFrame = cv2.resize(frame2,(int(width/2),int(height/2)));
            cv2.imshow('display',showFrame)
            cv2.setMouseCallback('display',centerCallback) #boundary callback
            key = cv2.waitKey(0)
            if key == 'x':
               resetThis = True
               break
            if key == 'q':
               exit()

            leftX, leftY = arenaX, arenaY;

            frame2 = copy.copy(frame);
            cv2.putText(frame2, 'Select RIGHT boundary', (50,50), cv2.FONT_HERSHEY_PLAIN, 2, 255)
            showFrame = cv2.resize(frame2,(int(width/2),int(height/2)));
            cv2.imshow('display',showFrame)
            cv2.setMouseCallback('display',centerCallback) #boundary callback
            key = cv2.waitKey(0)
            rightX, rightY = arenaX, arenaY;
            if key == 'x':
               resetThis = True
               break
            if key == 'q':
               exit()


            # show detected circle
            radius = np.sqrt((leftX - rightX)**2 + (leftY - rightY)**2)/2;
            frame2 = copy.copy(frame);
            cv2.putText(frame2, 'Here is the arena...I hope that is correct!', (50,50), cv2.FONT_HERSHEY_PLAIN, 2, 255)
            showFrame = cv2.resize(frame2,(int(width/2),int(height/2)));
            cv2.circle(showFrame,(centerX,centerY),int(radius),(0, 255, 0), 3)
            cv2.imshow('display',showFrame)
            key = cv2.waitKey(0)


            if key == 'x':
               resetThis = True
               break
            if key == 'q':
               exit()

         # step four: choose microphone centers, in order. show example of the proper ordering, then display best fit centers
         resetThis = True
         while resetThis == True:
            resetThis = False
            for micNum in list(range(16)):
               frame2 = copy.copy(frame);
               cv2.putText(frame2, 'Choose microphone number ' + str(micNum), (50,50), cv2.FONT_HERSHEY_PLAIN, 2, 255)
               showFrame = cv2.resize(frame2,(int(width/2),int(height/2)));
               cv2.imshow('display',showFrame)
               cv2.setMouseCallback('display',centerCallback)
               key = cv2.waitKey(0)
               if key == 'x':
                  micCenters = []
                  resetThis = True
                  break
               if key == 'q':
                  exit()

               micCenters.append((arenaX*2,arenaY*2))

         # we want to show the estimated position of the microphones to make sure we have them all...
         # maybe this should be shown in the 'align microphone' phase

         # step five: choose area that shows blinky light
         # choose blinky light area
         resetThis = True
         while resetThis == True:
            resetThis = False
            frame2 = copy.copy(frame);
            cv2.putText(frame2, 'Select blinky light', (50,50), cv2.FONT_HERSHEY_PLAIN, 2, 255)
            showFrame = cv2.resize(frame2,(int(width/2),int(height/2)));
            cv2.imshow('display',showFrame)
            cv2.setMouseCallback('display',areaCallback) #boundary callback
            key = cv2.waitKey(0)

            if key == 'x':
                  resetThis = True
                  break
            if key == 'q':
               exit()

      else:
         pass

   print('saving ' + fileName[:-4] + '_annotated2.mat')
   sio.savemat(fileName[:-4] + '_annotated2.mat',{'width':width,'height':height,'arenaCenter':(centerX*2,centerY*2),'arenaRadius':radius*2,
                                                                  'LEDcenter':(rectCenterX*2,rectCenterY*2),'LEDradius':rectRadius*2,'micCenters':micCenters})


cv2.destroyAllWindows()