latest changes

etrack/io/nixtrack_data.py

@@ -10,12 +10,14 @@ from IPython import embed
 
 class NixtrackData(object):
     
-    def __init__(self, results_file, crop=(0, 0)) -> None:
+    def __init__(self, filename, crop=(0, 0)) -> None:
         """
         If the video data was cropped before tracking and the tracked positions are with respect to the cropped images, we may want to correct for this to convert the data back to absolute positions in the video frame.
 
         Parameters
         ----------
+        filename : str
+            full filename
         crop : 2-tuple
             tuple of (xoffset, yoffset)
 
@@ -23,13 +25,13 @@ class NixtrackData(object):
         ------
         ValueError if crop value is not a 2-tuple
         """
-        if not os.path.exists(results_file):
-            raise ValueError("File %s does not exist!" % results_file)
+        if not os.path.exists(filename):
+            raise ValueError("File %s does not exist!" % filename)
         if not isinstance(crop, tuple) or len(crop) < 2:
             raise ValueError("Cropping info must be a 2-tuple of (x, y)")
-        self._file_name = results_file
+        self._file_name = filename
         self._crop = crop
-        self._dataset = nt.Dataset(self._file_name, nt.util.FileMode.ReadOnly)
+        self._dataset = nt.Dataset(self._file_name)
         if not self._dataset.is_open:
             raise ValueError(f"An error occurred opening file {self._file_name}! File is not open!")
 
@@ -50,14 +52,20 @@ class NixtrackData(object):
     def tracks(self):
         return self._dataset.tracks
 
-    def track(self, track=None, bodypart=0):
+    def track(self, bodypart=0, fps=None):
         if isinstance(bodypart, nb.Number):
             bp = self.bodyparts[bodypart]
-        elif isinstance(bodypart, str) and bodypart in self.bodyparts:
+        elif isinstance(bodypart, (str)) and bodypart in self.bodyparts:
             bp = bodypart
         else:
             raise ValueError(f"Body part {bodypart} is not a tracked node!")
-        positions, time, iscore, nscore = self._dataset.positions(node=bp, axis_type=nt.AxisType.Time)
-        embed()
+        if fps is None:
+            fps = self._dataset.fps
 
-        return TrackingData(positions[:, 0], positions[:, 1], time, nscore, bp, fps=self._dataset.fps)
+        positions, time, _, nscore = self._dataset.positions(node=bp, axis_type=nt.AxisType.Time)
+        valid = ~np.isnan(positions[:, 0])
+        positions = positions[valid,:]
+        time = time[valid]
+        score = nscore[valid]
+
+        return TrackingData(positions[:, 0], positions[:, 1], time, score, bp, fps=fps)