DeepLabCut · ehsainit · Jun 10, 2020 · Jun 11, 2020 · Jun 12, 2020 · Jun 12, 2020
diff --git a/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py
@@ -124,8 +124,11 @@ def strip_cropped_image_name(path):
     # ATTENTION: order has to be multibodyparts, then uniquebodyparts (for indexing)
     print("Utilizing the following graph:", partaffinityfield_graph)
     num_limbs = len(partaffinityfield_graph)
-    partaffinityfield_predict = True
-
+    # SO tracker one bp per ind => npo pafs
+    if len(partaffinityfield_graph) == 0:
+        partaffinityfield_predict = False
+    else:
+        partaffinityfield_predict = True
     # Loading the encoder (if necessary downloading from TF)
     dlcparent_path = auxiliaryfunctions.get_deeplabcut_path()
     defaultconfigfile = os.path.join(dlcparent_path, "pose_cfg.yaml")

diff --git a/deeplabcut/gui/analyze_videos.py b/deeplabcut/gui/analyze_videos.py
@@ -178,7 +178,7 @@ def __init__(self, parent, gui_size, cfg):
                 self, label="Specify the Tracker Method (you can try each)"
             )
             tracker_text_boxsizer = wx.StaticBoxSizer(tracker_text, wx.VERTICAL)
-            trackertypes = ["skeleton", "box"]
+            trackertypes = ["skeleton", "box", "single_object"]
             self.trackertypes = wx.ComboBox(
                 self, choices=trackertypes, style=wx.CB_READONLY
             )

diff --git a/deeplabcut/gui/create_videos.py b/deeplabcut/gui/create_videos.py
@@ -207,7 +207,7 @@ def __init__(self, parent, gui_size, cfg):
         if self.cfg.get("multianimalproject", False):
             tracker_text = wx.StaticBox(self, label="Specify the Tracker Method!")
             tracker_text_boxsizer = wx.StaticBoxSizer(tracker_text, wx.VERTICAL)
-            trackertypes = ["skeleton", "box"]
+            trackertypes = ["skeleton", "box", "single_object"]
             self.trackertypes = wx.ComboBox(
                 self, choices=trackertypes, style=wx.CB_READONLY
             )

diff --git a/deeplabcut/pose_estimation_tensorflow/lib/single_object_tracker.py b/deeplabcut/pose_estimation_tensorflow/lib/single_object_tracker.py
@@ -0,0 +1,120 @@
+"""
+DeepLabCut 2.2 Toolbox (deeplabcut.org)
+https://github.com/DeepLabCut/DeepLabCut
+
+© DeepLabCut authors
+Please see AUTHORS for contributors.
+https://github.com/AlexEMG/DeepLabCut/blob/master/AUTHORS
+
+Licensed under GNU Lesser General Public License v3.0
+
+This module was contributed by Tabet Ehsainieh - https://github.com/ehsainit
+"""
+
+from deeplabcut.pose_estimation_tensorflow.lib.trackingutils import *
+
+
+class TrackByDetectionTracker:
+    def __init__(self, max_frames_to_skip, jointsnum):
+        self.max_frames_to_skip = max_frames_to_skip
+        self.id_count = 0
+        self.tracks = []
+        self.iter = 0
+        self.det = []
+        self.jointnum = jointsnum
+
+    def cost_metric(self, detections, N, M):
+        cost_matrix = np.zeros(shape=(N, M))  # Cost matrix
+        for i in range(len(self.tracks)):
+            for j in range(len(detections)):
+                z = convert_bbox_to_z(detections[j])[:2]
+                gating_dist = self.tracks[i].KF.mahalanobis_dist(z)
+                cost_matrix[i][j] = gating_dist
+        return cost_matrix
+
+    def track(self, bbs):
+        self.iter += 1
+        detection_bbs = bbs
+
+        if len(self.tracks) == 0:
+            for i in range(len(detection_bbs)):
+                track = Track(detection_bbs[i], self.id_count)
+                self.id_count += 1
+                self.tracks.append(track)
+
+        N = len(self.tracks)
+        M = len(detection_bbs)
+        cost_matrix = self.cost_metric(detection_bbs, N, M)
+        rows, cols = linear_sum_assignment(cost_matrix)
+        matches = self._matching(rows, cols, bbs, N)
+        # either predict or use detection bbs to track
+        for i in range(len(matches)):
+            # would be better if delete from the matches upon deleting tracks
+            if i < len(self.tracks):
+                if matches[i] is not None:
+                    self.tracks[i].skipped_frames = 0
+                    self.tracks[i].predict(self.iter, detection_bbs[matches[i]])
+                else:
+                    self.tracks[i].predict(it=self.iter)
+        states = []
+        for t, track in enumerate(self.tracks):
+            if t < len(matches) and matches[t] is not None:
+                # new assigned tracks in the n frames will be fist tracked in the n+1 frame - probably need fixing
+                states.append(np.concatenate((track.prediction, [track.track_id, matches[t]])).reshape(1, -1)[0])
+        if len(states) > 0:
+            return np.stack(states)
+        else:
+            return np.empty((0, self.jointnum * 2 + 2))
+
+    def _matching(self, rows, cols, detection_bbs, N):
+
+        # assign detections as matches according to the linear assignment solution
+        matches = [None] * N
+        for i in range(len(rows)):
+            matches[rows[i]] = cols[i]
+
+        # validate matches :  if the iou are unreasonable unmatch, if linear assignment didnt find
+        # a solution increase the number of skipped frames
+        for i in range(len(matches)):
+            if matches[i] is not None:
+                _iou = iou(self.tracks[i].prediction, detection_bbs[matches[i]])
+                if _iou > 9:
+                    print("unmatched found")
+                    matches[i] = None
+            else:
+                self.tracks[i].skipped_frames += 1
+        # handle tracks with skipped frame > max_frame_to_skip -> assign death
+        del_tracks = []
+        for i in range(len(self.tracks)):
+            if self.tracks[i].skipped_frames > self.max_frames_to_skip:
+                del_tracks.append(i)
+
+        if len(del_tracks) > 0:
+            for i in del_tracks:
+                del self.tracks[i]
+
+        # for detections with no matches will be considered as births - new tracks
+        for i in range(len(detection_bbs)):
+            if i not in matches:
+                track = Track(detection_bbs[i], self.id_count)
+                self.id_count += 1
+                self.tracks.append(track)
+        return matches
+
+
+class Track:
+    def __init__(self, corr_bb, trackId):
+        self.track_id = trackId
+        self.KF = KalmanBoxTracker(corr_bb)
+        self.prediction = corr_bb
+        self.trace = dict()  # trace path
+        self.skipped_frames = 0
+
+    def predict(self, it, detection=None):
+        if detection is not None:
+            self.KF.update(detection)
+            self.KF.predict()
+            self.prediction = self.KF.get_state()[0]
+        self.trace[it] = self.prediction # probably will need only the center
+        self.KF.predict()
+        return self.KF.get_state()[0]
diff --git a/deeplabcut/pose_estimation_tensorflow/lib/trackingutils.py b/deeplabcut/pose_estimation_tensorflow/lib/trackingutils.py
@@ -166,6 +166,17 @@ def get_state(self):
         Returns the current bounding box estimate.
         """
         return convert_x_to_bbox(self.kf.x)
+
+    def mahalanobis_dist(self, z):
+        """
+           Compute the mahalanobis distance
+        """
+        x_hat = np.dot(self.kf.H[:2,:2], self.kf.x[:2])
+        y = np.subtract(z, x_hat)
+        S = self.kf.R[:2,:2] + np.dot(self.kf.H[:2,:2], np.dot(self.kf.P[:2,:2], self.kf.H.T[:2,:2]))
+        Inv_S = np.linalg.inv(S)
+        d = np.dot(y.T, np.dot(Inv_S, y))
+        return np.sqrt(d)
 
 
 class SkeletonTracker:

diff --git a/deeplabcut/pose_estimation_tensorflow/predict_videos.py b/deeplabcut/pose_estimation_tensorflow/predict_videos.py
@@ -1224,11 +1224,12 @@ def convert_detections2tracklets(
 
     """
     from deeplabcut.pose_estimation_tensorflow.lib import inferenceutils, trackingutils
+    from deeplabcut.pose_estimation_tensorflow.lib.single_object_tracker import TrackByDetectionTracker
     from deeplabcut.utils import auxfun_multianimal
     from easydict import EasyDict as edict
     import pickle
 
-    if track_method not in ("box", "skeleton"):
+    if track_method not in ("box", "skeleton", "single_object"):
         raise ValueError(
             "Invalid tracking method. Only `box` and `skeleton` are currently supported."
         )
@@ -1344,7 +1345,14 @@ def convert_detections2tracklets(
             vname = Path(video).stem
             dataname = os.path.join(videofolder, vname + DLCscorer + ".h5")
             data, metadata = auxfun_multianimal.LoadFullMultiAnimalData(dataname)
-            method = "sk" if track_method == "skeleton" else "bx"
+            if track_method == "skeleton":
+                method = "sk"
+            elif track_method == "single_object":
+                method = "so"
+            elif track_method == "box":
+                method = "bx"
+            else:
+                raise ValueError
             trackname = dataname.split(".h5")[0] + f"_{method}.pickle"
             trackname = trackname.replace(videofolder, destfolder)
             if (
@@ -1419,6 +1427,8 @@ def convert_detections2tracklets(
 
                 if track_method == "box":
                     mot_tracker = trackingutils.Sort(inferencecfg)
+                elif track_method == "single_object":
+                    mot_tracker = TrackByDetectionTracker(10, numjoints)
                 else:
                     mot_tracker = trackingutils.SORT(
                         numjoints,
@@ -1448,6 +1458,10 @@ def convert_detections2tracklets(
                             inferencecfg, animals, 0
                         )  # TODO: get cropping parameters and utilize!
                         trackers = mot_tracker.update(bb)
+                    elif track_method == "single_object":
+                        bb = inferenceutils.individual2boundingbox(inferencecfg, animals,
+                                                                   0)  # TODO: get cropping parameters and utilize!
+                        trackers = mot_tracker.track(bb)
                     else:
                         temp = [arr.reshape((-1, 3))[:, :2] for arr in animals]
                         trackers = mot_tracker.track(temp)

diff --git a/deeplabcut/utils/auxfun_multianimal.py b/deeplabcut/utils/auxfun_multianimal.py
@@ -57,25 +57,28 @@ def getpafgraph(cfg, printnames=True):
     # CHECKS if each bpt is connected to at least one other bpt
     # TODO: check that there is a path leading from each (multi)bpt to each other (multi)bpt!
     connected = set()
-    partaffinityfield_graph = []
-    for link in cfg["skeleton"]:
-        if link[0] in bodypartnames and link[1] in bodypartnames:
-            bp1 = int(lookupdict[link[0]])
-            bp2 = int(lookupdict[link[1]])
-            connected.add(bp1)
-            connected.add(bp2)
-            partaffinityfield_graph.append([bp1, bp2])
-        else:
-            print("Attention, parts do not exist!", link)
-
-    unconnected = set(range(len(multianimalbodyparts))).difference(connected)
-    if unconnected:
-        raise ValueError(
-            f'Unconnected {", ".join(multianimalbodyparts[i] for i  in unconnected)}. '
-            f"For multi-animal projects, all multianimalbodyparts should be connected. "
-            f"Ideally there should be at least one (multinode) path from each multianimalbodyparts to each other multianimalbodyparts. "
-            f"Please verify the skeleton in the config.yaml."
-        )
+    if len(cfg["multianimalbodyparts"]) == 1:
+        partaffinityfield_graph = []
+    else:
+        partaffinityfield_graph = []
+        for link in cfg["skeleton"]:
+            if link[0] in bodypartnames and link[1] in bodypartnames:
+                bp1 = int(lookupdict[link[0]])
+                bp2 = int(lookupdict[link[1]])
+                connected.add(bp1)
+                connected.add(bp2)
+                partaffinityfield_graph.append([bp1, bp2])
+            else:
+                print("Attention, parts do not exist!", link)
+
+        unconnected = set(range(len(multianimalbodyparts))).difference(connected)
+        if unconnected:
+            raise ValueError(
+                f'Unconnected {", ".join(multianimalbodyparts[i] for i  in unconnected)}. '
+                f"For multi-animal projects, all multianimalbodyparts should be connected. "
+                f"Ideally there should be at least one (multinode) path from each multianimalbodyparts to each other multianimalbodyparts. "
+                f"Please verify the skeleton in the config.yaml."
+            )
 
     if printnames:
         graph2names(cfg, partaffinityfield_graph)

diff --git a/deeplabcut/utils/auxiliaryfunctions.py b/deeplabcut/utils/auxiliaryfunctions.py
@@ -753,6 +753,8 @@ def load_detection_data(video, scorer, track_method):
         tracker = "sk"
     elif track_method == "box":
         tracker = "bx"
+    elif track_method == "single_object":
+        tracker = "so"
     else:
         raise ValueError(f"Unrecognized track_method={track_method}")