From e116c61e0050723bf18b7233c63d4c05ae3c9089 Mon Sep 17 00:00:00 2001
From: Arya Massarat <aryam7@gmail.com>
Date: Fri, 12 Apr 2019 14:29:35 -0700
Subject: [PATCH] develop a new criterion for determining ant direction. see
 #18

---
 track.py | 60 ++++++++++++++++++++++----------------------------------
 1 file changed, 23 insertions(+), 37 deletions(-)

diff --git a/track.py b/track.py
index 233339a..34b090d 100755
--- a/track.py
+++ b/track.py
@@ -7,7 +7,7 @@
 import argparse
 
 
-def trackOneClip(vidPath, cushion, W, H, minBlob, vidExport, result_path):
+def trackOneClip(vidPath, W, H, minBlob, vidExport, result_path):
     eng = matlab.engine.start_matlab()
     # call the ant_tracking.m script and get the resulting dataframe
     # inputs:
@@ -17,7 +17,7 @@ def trackOneClip(vidPath, cushion, W, H, minBlob, vidExport, result_path):
     #   result_path - string, path to the directory in which to store result videos
     df = eng.ant_tracking(vidExport, minBlob, vidPath, result_path)
     if df:
-        track_result = np.array([["fName", "X", "Y"]])
+        track_result = np.array([["fName", "id", "X", "Y"]])
         # convert the dataframe to a np array
         # it should have five columns:
         #   x_pos, y_pos, width, height, ant_id
@@ -28,36 +28,18 @@ def trackOneClip(vidPath, cushion, W, H, minBlob, vidExport, result_path):
         for idnum in idL:
             # get tracks for this ant
             antTrack = df[df[:, 4] == idnum]
-            # get first position
-            begin = antTrack[0,:]
-            # get last position
-            end = antTrack[-1,:]
-            # get start and end coordinates
-            # and width and height of the bounding boxes
-            x0 = begin[0]
-            y0 = begin[1]
-            x1 = end[0]
-            y1 = end[1]
-            width0 = begin[2]
-            height0 = begin[3]
-            width1 = end[2]
-            height1 = end[3]
-            # attempt to determine X and Y directions
-            directionX = "NA"
-            directionY = "NA"
-            if x0 <= cushion and x1+width1>=(W-cushion):
-                directionX = "LR"
-            elif (x0 + width0) >= (W-cushion) and x1 <= cushion:
-                directionX = "RL"
-            if  y0 <= cushion and (y1+height1)>= (H-cushion):
-                directionY = "TB"
-            elif  (y0+height0) >= (H-cushion) and y1<= cushion:
-                directionY = "BT"
+            # NOTE: x and y coords can be negative if kalman filter is predicting the ant after it passes out of frame
+            x0 = antTrack[1:,0]
+            x1 = antTrack[:-1,0]
+            directionX = sum(x0-x1)
+            y0 = antTrack[1:,1]
+            y1 = antTrack[:-1,1]
+            directionY = sum(y0-y1)
             # save results in np array so that we can return them soon
-            track_result = np.append(track_result,  [[vidPath, directionX, directionY]], axis=0)
+            track_result = np.append(track_result, [[vidPath, idnum, directionX, directionY]], axis=0)
         # return the data without its header
-        return track_result[1:,]
-    return np.array([])
+        return track_result[1:,], df
+    return np.array([]), np.array([])
 
 def main():
     arg_parser = argparse.ArgumentParser()
@@ -76,11 +58,11 @@ def main():
                             required = True,
                             type=str,
                             help='path of array of results to output')
-    arg_parser.add_argument('--c',
-                            dest = 'cushion',
-                            type=int,
-                            default=10,
-                            help='number of pixels as padding (default = 10)')
+    arg_parser.add_argument('--rr',
+                            dest = 'raw_results',
+                            type=str,
+                            default=None,
+                            help='path of array of raw results to output (default = None)')
     arg_parser.add_argument('--m',
                             dest = 'minBlob',
                             type=int,
@@ -91,20 +73,24 @@ def main():
                             type=bool,
                             default=False,
                             help='do we export result video (default = False)')
+
     args = arg_parser.parse_args()
 
     # track ants in each of the cropped videos
-    result_array = np.array([["fName", "X", "Y"]])
+    result_array = np.array([["fName", "id", "X", "Y"]])
     print "Tracking ants in " + args.cropVid
     # get height and width of video
     H, W = findVideoMetada(args.cropVid)
     # call matlab to track ants in a single cropped video
-    track_result = trackOneClip(args.cropVid, args.cushion, W, H, args.minBlob, args.export, args.result_path)
+    track_result, raw_results = trackOneClip(args.cropVid, W, H, args.minBlob, args.export, args.result_path)
     # keep track of the tracking results in a np array
     if track_result.size:
         result_array = np.concatenate((result_array, track_result), axis=0)
     # save the tracking results to disk
     np.savetxt(args.result, result_array, delimiter= ',', fmt='%s')
+    # save the raw results to disk
+    if args.raw_results is not None:
+        np.savetxt(args.raw_results, raw_results, delimiter=',', fmt='%s')
 
 if __name__== "__main__":
     main()