[WIP] Add tracklet properties

btrack/btypes.py

@@ -19,6 +19,7 @@
 
 import ctypes
 from collections import OrderedDict
+from typing import Dict
 
 import numpy as np
 
@@ -357,6 +358,7 @@ def __init__(
         self.ID = ID
         self._data = data
         self._kalman = None
+        self._properties = {}
 
         self.root = None
         self.parent = parent
@@ -383,47 +385,67 @@ def _repr_html_(self):
             )
 
     @property
-    def x(self):
+    def properties(self) -> Dict[str, np.ndarray]:
+        return self._properties
+
+    @properties.setter
+    def properties(self, properties: Dict[str, np.ndarray]):
+        """Store properties associated with this Tracklet."""
+
+        # validate the track properties
+        for k, v in properties.items():
+            if len(v) != len(self):
+                raise ValueError(
+                    'The number of properties and track objects must be equal.'
+                )
+            # ensure the property values are a numpy array
+            if type(v) != np.ndarray:
+                properties[k] = np.asarray(v)
+
+        self._properties = properties
+
+    @property
+    def x(self) -> list:
         return [o.x for o in self._data]
 
     @property
-    def y(self):
+    def y(self) -> list:
         return [o.y for o in self._data]
 
     @property
-    def z(self):
+    def z(self) -> list:
         return [o.z for o in self._data]
 
     @property
-    def t(self):
+    def t(self) -> list:
         return [o.t for o in self._data]
 
     @property
-    def dummy(self):
+    def dummy(self) -> list:
         return [o.dummy for o in self._data]
 
     @property
-    def refs(self):
+    def refs(self) -> list:
         return [o.ID for o in self._data]
 
     @property
-    def start(self):
+    def start(self) -> list:
         return self.t[0]
 
     @property
-    def stop(self):
+    def stop(self) -> list:
         return self.t[-1]
 
     @property
-    def label(self):
+    def label(self) -> list:
         return [o.state.name for o in self._data]
 
     @property
-    def state(self):
+    def state(self) -> list:
         return [o.state.value for o in self._data]
 
     @property
-    def softmax(self):
+    def softmax(self) -> list:
         return [o.probability for o in self._data]
 
     @property

tests/test_tracklets.py

@@ -0,0 +1,88 @@
+import numpy as np
+import pytest
+
+import btrack
+
+
+def _create_test_object():
+    """Create a test object."""
+    data = {
+        "ID": np.random.randint(0, 1000),
+        "x": np.random.uniform(0.0, 1000.0),
+        "y": np.random.uniform(0.0, 1000.0),
+        "z": np.random.uniform(0.0, 1000.0),
+        "t": np.random.randint(0, 1000),
+        "dummy": False,
+        "states": 5,
+        "label": 0,
+        "prob": 0.5,
+    }
+
+    obj = btrack.btypes.PyTrackObject().from_dict(data)
+    return obj, data
+
+
+def _create_test_tracklet(track_len: int):
+    """Create a test track."""
+    data = [_create_test_object()[0] for i in range(track_len)]
+    track_ID = np.random.randint(0, 1000)
+    tracklet = btrack.btypes.Tracklet(track_ID, data)
+    return tracklet, data, track_ID
+
+
+def _create_test_properties(track_len: int):
+    properties = {
+        "speed": np.random.uniform(0.0, 1.0, size=(track_len,)),
+        "circularity": np.random.uniform(0.0, 1.0, size=(track_len,)),
+        "reporter": np.random.uniform(0.0, 1.0, size=(track_len,)),
+    }
+    return properties
+
+
+def test_object():
+    """Test that an object is correctly instantiated, and that the stored
+    data matches the data used for creation."""
+    obj, data = _create_test_object()
+    for k, v in data.items():
+        assert getattr(obj, k) == v
+
+
+@pytest.mark.parametrize("track_len", [0, 1, 10, 100, 1000])
+def test_tracklet(track_len: int):
+    """Test that a track is correctly instantiated, and that the stored
+    data matches the data used for creation."""
+    tracklet, data, track_ID = _create_test_tracklet(track_len)
+    assert len(tracklet) == len(data)
+
+    # now check that the track data is correct
+    assert tracklet.ID == track_ID
+    fields = ["x", "y", "z", "t"]
+    for field in fields:
+        obj_data = [getattr(obj, field) for obj in data]
+        np.testing.assert_equal(obj_data, getattr(tracklet, field))
+
+
+@pytest.mark.parametrize("track_len", [0, 1, 10, 100, 1000])
+def test_tracklet_properties(track_len: int):
+    """Test that a track is correctly instantiated, and that the stored
+    properties match the data used for creation."""
+    tracklet, data, track_ID = _create_test_tracklet(track_len)
+    properties = _create_test_properties(track_len)
+
+    tracklet.properties = properties
+    t_properties = tracklet.properties
+
+    for k, v in properties.items():
+        np.testing.assert_equal(t_properties[k], v)
+
+
+def test_malformed_tracklet_properties():
+    """Test for malformed properties by truncating the first property."""
+    track_len = 10
+    tracklet, data, track_ID = _create_test_tracklet(track_len)
+    properties = _create_test_properties(track_len)
+    first_key = list(properties.keys())[0]
+    properties[first_key] = properties[first_key][: track_len - 2]
+
+    with pytest.raises(ValueError):
+        tracklet.properties = properties