diff --git a/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py index 9c620c1093..df17b6849f 100644 --- 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 index 75a1286ff8..0d492621e9 100644 --- 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 index 7fde61fb0b..f7d7301909 100644 --- 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 new file mode 100644 index 0000000000..3fc9ac2215 --- /dev/null +++ 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 index 3969cebb44..dc7b8dc018 100644 --- 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 index ac28296e14..e70f54f5b6 100755 --- 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 index 1d4a309869..7f9ee1d922 100644 --- 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 index a41fe8240a..b840c9af9a 100644 --- 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}")