diff --git a/deeplabcut/lmot/__init__.py b/deeplabcut/lmot/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/deeplabcut/lmot/assets.py b/deeplabcut/lmot/assets.py new file mode 100644 index 0000000000..3d3784950b --- /dev/null +++ b/deeplabcut/lmot/assets.py @@ -0,0 +1,26 @@ +import random + +import cv2 + + +def get_random_color(n): + ''' generate rgb using a list comprehension ''' + r = int(random.random() * 256) + g = int(random.random() * 256) + b = int(random.random() * 256) + step = 256 / n + for i in range(n): + r += step + g += step + b += step + r = int(r) % 256 + g = int(g) % 256 + b = int(b) % 256 + return r, g, b + + +def bbox(frame, track): + x = track.prediction[0] + y = track.prediction[1] + return cv2.rectangle(frame, (int(x + track.width + 1), int(y + track.height + 1)), + (int(x - track.width + 1), int(y - track.height + 1)), track.color, 2) diff --git a/deeplabcut/lmot/detection.py b/deeplabcut/lmot/detection.py new file mode 100644 index 0000000000..8a9fb4d7e0 --- /dev/null +++ b/deeplabcut/lmot/detection.py @@ -0,0 +1,5 @@ +class Detection: + def __init__(self, coordinates, height, width): + self.coordinates = coordinates + self.height = height + self.width = width diff --git a/deeplabcut/lmot/extract_local_maxima.py b/deeplabcut/lmot/extract_local_maxima.py new file mode 100644 index 0000000000..0fa3bf72a3 --- /dev/null +++ b/deeplabcut/lmot/extract_local_maxima.py @@ -0,0 +1,36 @@ +import numpy as np +import scipy.ndimage as ndimage +import scipy.ndimage.filters as filters + +from deeplabcut.lmot.detection import Detection + + +def extract_locmaxima(scmap, locref, neighborhood_size=5, threshold=0.99999): + # credit: https://stackoverflow.com/questions/9111711/get-coordinates-of-local-maxima-in-2d-array-above-certain-value + data_max = filters.maximum_filter(scmap, neighborhood_size) + maxima = (scmap == data_max) + data_min = filters.minimum_filter(scmap, neighborhood_size) + diff = ((data_max - data_min) > threshold) + maxima[diff == 0] = 0 + labeled, num_objects = ndimage.label(maxima) + possible_objects = ndimage.find_objects(labeled) + detections = [] + for dy, dx in possible_objects: + height, width, center = compute_bbox(dy, dx, locref) + measurement = Detection(center, height, width) + detections.append(measurement) + return detections + + +def compute_bbox(dy, dx, locref): + startloc = extract_point_from_nn((dy.start, dx.start), locref) + width = extract_point_from_nn((dy.start, dx.stop), locref) - extract_point_from_nn((dy.start, dx.start), locref) + height = extract_point_from_nn((dy.stop, dx.start), locref) - extract_point_from_nn((dy.start, dx.start), locref) + return abs(height[1]), abs(width[0]), startloc + + +def extract_point_from_nn(loc, locref): + offset = np.array(locref[loc])[::-1] + pos_f8 = (np.array(loc).astype('float') * 8.0 + 0.5 * 8.0 + offset) + pose = np.array([pos_f8[1], pos_f8[0], 0.0, 0.0]) + return pose diff --git a/deeplabcut/lmot/kalman_filter.py b/deeplabcut/lmot/kalman_filter.py new file mode 100644 index 0000000000..b3f84f65eb --- /dev/null +++ b/deeplabcut/lmot/kalman_filter.py @@ -0,0 +1,101 @@ +from copy import deepcopy + +import numpy as np + + +class KalmanFilter: + + def __init__(self, x): + self.dt = 0.005 # delta time + + self.H = np.array([[1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 0, 0], + [0, 0, 0, 0]]) # matrix in observation equations + + self.x = x # initial state estimate + self.Hmd = np.array([[1, 0], + [0, 1]]) + self.P = np.array([[3.0, 0, 0, 0], + [0, 3.0, 0, 0], + [0, 0, 3.0, 0], + [0, 0, 0, 3.0]]) # covariance matrix + self.F = np.array([[1, 0, self.dt, 0], + [0, 1, 0, self.dt], + [0, 0, 1, 0], + [0, 0, 0, 1]]) # state transition mat + self.Q = 10000 * np.array( + [[(self.dt ** 4) / 4, 0, (self.dt ** 3) / 2, 0], + [0, (self.dt ** 4) / 4, 0, (self.dt ** 3) / 2], + [(self.dt ** 3) / 2, 0, self.dt ** 2, 0], + [0, (self.dt ** 3) / 2, 0, self.dt ** 2]]) + # process noise matrix + self.R = 0.1 * np.eye(self.x.shape[0]) # observation noise matrix + self.G = np.array([[self.dt ** 2 / 2], [self.dt ** 2 / 2], [self.dt], [self.dt]]) + + self.z = np.array([[None] * 4]).T + # these will always be a copy of x,P after predict() is called + self.x_prior = self.x.copy() + self.P_prior = self.P.copy() + + # these will always be a copy of x,P after update() is called + self.x_post = self.x.copy() + self.P_post = self.P.copy() + # read only - to keep track and further calculations + self.S = None # innovation (pre-fit residual) covariance + self.y = None # measurement pre-fit residual + self.K = None # Optimal Gain + self.Inv_S = None + + def predict(self): + # Predicted state estimate + self.x = np.dot(self.F, self.x) + self.P = np.dot(self.F, np.dot(self.P, self.F.T)) + self.Q + # save prior + self.x_prior = self.x.copy() + self.P_prior = self.P.copy() + return np.round(self.x) + + def update(self, z): + y = z - np.dot(self.H, self.x) + S = self.R + np.dot(self.H, np.dot(self.P, self.H.T)) + K = np.dot(self.P, np.dot(self.H.T, np.linalg.inv(S))) + self.x = self.x + np.dot(K, y) + self.P = np.dot((np.identity(4) - np.dot(K, self.H)), self.P) + self.z = deepcopy(z) + self.x_post = self.x.copy() + self.P_post = self.P.copy() + # read only - to keep track for further calculations + self.y = y + self.S = S + self.Inv_S = np.linalg.inv(S) + self.K = K + return np.round(self.x) + + def mahalanobis_dist(self, z): + """ + Compute the mahalanobis distance + """ + x_hat = np.dot(self.Hmd, self.x[:2]) + y = z[:2] - x_hat + S = self.R + np.dot(self.H, np.dot(self.P, self.H.T)) + S = S[: 2:] + S = S[:, :2] + Inv_S = np.linalg.inv(S) + d = np.dot(y.T, np.dot(Inv_S, y)) + return np.sqrt(d) + + def GetPreRes(self): + if self.y is not None: + return self.y + else: + return 0 + + def estimateCov(self): + pass + + def get_xPrior(self): + return self.x_prior + + def _update(self): + pass diff --git a/deeplabcut/lmot/linear_sum_assignment.py b/deeplabcut/lmot/linear_sum_assignment.py new file mode 100644 index 0000000000..16bf2ab766 --- /dev/null +++ b/deeplabcut/lmot/linear_sum_assignment.py @@ -0,0 +1,287 @@ +# Hungarian algorithm (Kuhn-Munkres) for solving the linear sum assignment +# problem. Taken from scikit-learn. Based on original code by Brian Clapper, +# adapted to NumPy by Gael Varoquaux. +# Further improvements by Ben Root, Vlad Niculae and Lars Buitinck. +# +# Copyright (c) 2008 Brian M. Clapper , Gael Varoquaux +# Author: Brian M. Clapper, Gael Varoquaux +# License: 3-clause BSD + +import numpy as np + + +def linear_sum_assignment(cost_matrix): + """Solve the linear sum assignment problem. + + The linear sum assignment problem is also known as minimum weight matching + in bipartite graphs. A problem instance is described by a matrix C, where + each C[i,j] is the cost of matching vertex i of the first partite set + (a "worker") and vertex j of the second set (a "job"). The goal is to find + a complete assignment of workers to jobs of minimal cost. + + Formally, let X be a boolean matrix where :math:`X[i,j] = 1` iff row i is + assigned to column j. Then the optimal assignment has cost + + .. math:: + \\min \\sum_i \\sum_j C_{i,j} X_{i,j} + + s.t. each row is assignment to at most one column, and each column to at + most one row. + + This function can also solve a generalization of the classic assignment + problem where the cost matrix is rectangular. If it has more rows than + columns, then not every row needs to be assigned to a column, and vice + versa. + + The method used is the Hungarian algorithm, also known as the Munkres or + Kuhn-Munkres algorithm. + + Parameters + ---------- + cost_matrix : array + The cost matrix of the bipartite graph. + + Returns + ------- + row_ind, col_ind : array + An array of row indices and one of corresponding column indices giving + the optimal assignment. The cost of the assignment can be computed + as ``cost_matrix[row_ind, col_ind].sum()``. The row indices will be + sorted; in the case of a square cost matrix they will be equal to + ``numpy.arange(cost_matrix.shape[0])``. + + Notes + ----- + .. versionadded:: 0.17.0 + + Examples + -------- + >>> cost = np.array([[4, 1, 3], [2, 0, 5], [3, 2, 2]]) + >>> from scipy.optimize import linear_sum_assignment + >>> row_ind, col_ind = linear_sum_assignment(cost) + >>> col_ind + array([1, 0, 2]) + >>> cost[row_ind, col_ind].sum() + 5 + + References + ---------- + 1. http://csclab.murraystate.edu/bob.pilgrim/445/munkres.html + + 2. Harold W. Kuhn. The Hungarian Method for the assignment problem. + *Naval Research Logistics Quarterly*, 2:83-97, 1955. + + 3. Harold W. Kuhn. Variants of the Hungarian method for assignment + problems. *Naval Research Logistics Quarterly*, 3: 253-258, 1956. + + 4. Munkres, J. Algorithms for the Assignment and Transportation Problems. + *J. SIAM*, 5(1):32-38, March, 1957. + + 5. https://en.wikipedia.org/wiki/Hungarian_algorithm + """ + cost_matrix = np.asarray(cost_matrix) + if len(cost_matrix.shape) != 2: + raise ValueError("expected a matrix (2-d array), got a %r array" + % (cost_matrix.shape,)) + + if not (np.issubdtype(cost_matrix.dtype, np.number) or + cost_matrix.dtype == np.dtype(np.bool)): + raise ValueError("expected a matrix containing numerical entries, got %s" + % (cost_matrix.dtype,)) + + if np.any(np.isinf(cost_matrix) | np.isnan(cost_matrix)): + raise ValueError("matrix contains invalid numeric entries") + + if cost_matrix.dtype == np.dtype(np.bool): + cost_matrix = cost_matrix.astype(np.int) + + # The algorithm expects more columns than rows in the cost matrix. + if cost_matrix.shape[1] < cost_matrix.shape[0]: + cost_matrix = cost_matrix.T + transposed = True + else: + transposed = False + + state = _Hungary(cost_matrix) + + # No need to bother with assignments if one of the dimensions + # of the cost matrix is zero-length. + step = None if 0 in cost_matrix.shape else _step1 + + while step is not None: + step = step(state) + + if transposed: + marked = state.marked.T + else: + marked = state.marked + return np.where(marked == 1) + + +class _Hungary(object): + """State of the Hungarian algorithm. + + Parameters + ---------- + cost_matrix : 2D matrix + The cost matrix. Must have shape[1] >= shape[0]. + """ + + def __init__(self, cost_matrix): + self.C = cost_matrix.copy() + + n, m = self.C.shape + self.row_uncovered = np.ones(n, dtype=bool) + self.col_uncovered = np.ones(m, dtype=bool) + self.Z0_r = 0 + self.Z0_c = 0 + self.path = np.zeros((n + m, 2), dtype=int) + self.marked = np.zeros((n, m), dtype=int) + + def _clear_covers(self): + """Clear all covered matrix cells""" + self.row_uncovered[:] = True + self.col_uncovered[:] = True + + +# Individual steps of the algorithm follow, as a state machine: they return +# the next step to be taken (function to be called), if any. + +def _step1(state): + """Steps 1 and 2 in the Wikipedia page.""" + + # Step 1: For each row of the matrix, find the smallest element and + # subtract it from every element in its row. + state.C -= state.C.min(axis=1)[:, np.newaxis] + # Step 2: Find a zero (Z) in the resulting matrix. If there is no + # starred zero in its row or column, star Z. Repeat for each element + # in the matrix. + for i, j in zip(*np.where(state.C == 0)): + if state.col_uncovered[j] and state.row_uncovered[i]: + state.marked[i, j] = 1 + state.col_uncovered[j] = False + state.row_uncovered[i] = False + + state._clear_covers() + return _step3 + + +def _step3(state): + """ + Cover each column containing a starred zero. If n columns are covered, + the starred zeros describe a complete set of unique assignments. + In this case, Go to DONE, otherwise, Go to Step 4. + """ + marked = (state.marked == 1) + state.col_uncovered[np.any(marked, axis=0)] = False + + if marked.sum() < state.C.shape[0]: + return _step4 + + +def _step4(state): + """ + Find a noncovered zero and prime it. If there is no starred zero + in the row containing this primed zero, Go to Step 5. Otherwise, + cover this row and uncover the column containing the starred + zero. Continue in this manner until there are no uncovered zeros + left. Save the smallest uncovered value and Go to Step 6. + """ + # We convert to int as numpy operations are faster on int + C = (state.C == 0).astype(int) + covered_C = C * state.row_uncovered[:, np.newaxis] + covered_C *= np.asarray(state.col_uncovered, dtype=int) + n = state.C.shape[0] + m = state.C.shape[1] + + while True: + # Find an uncovered zero + row, col = np.unravel_index(np.argmax(covered_C), (n, m)) + if covered_C[row, col] == 0: + return _step6 + else: + state.marked[row, col] = 2 + # Find the first starred element in the row + star_col = np.argmax(state.marked[row] == 1) + if state.marked[row, star_col] != 1: + # Could not find one + state.Z0_r = row + state.Z0_c = col + return _step5 + else: + col = star_col + state.row_uncovered[row] = False + state.col_uncovered[col] = True + covered_C[:, col] = C[:, col] * ( + np.asarray(state.row_uncovered, dtype=int)) + covered_C[row] = 0 + + +def _step5(state): + """ + Construct a series of alternating primed and starred zeros as follows. + Let Z0 represent the uncovered primed zero found in Step 4. + Let Z1 denote the starred zero in the column of Z0 (if any). + Let Z2 denote the primed zero in the row of Z1 (there will always be one). + Continue until the series terminates at a primed zero that has no starred + zero in its column. Unstar each starred zero of the series, star each + primed zero of the series, erase all primes and uncover every line in the + matrix. Return to Step 3 + """ + count = 0 + path = state.path + path[count, 0] = state.Z0_r + path[count, 1] = state.Z0_c + + while True: + # Find the first starred element in the col defined by + # the path. + row = np.argmax(state.marked[:, path[count, 1]] == 1) + if state.marked[row, path[count, 1]] != 1: + # Could not find one + break + else: + count += 1 + path[count, 0] = row + path[count, 1] = path[count - 1, 1] + + # Find the first prime element in the row defined by the + # first path step + col = np.argmax(state.marked[path[count, 0]] == 2) + if state.marked[row, col] != 2: + col = -1 + count += 1 + path[count, 0] = path[count - 1, 0] + path[count, 1] = col + + # Convert paths + for i in range(count + 1): + if state.marked[path[i, 0], path[i, 1]] == 1: + state.marked[path[i, 0], path[i, 1]] = 0 + else: + state.marked[path[i, 0], path[i, 1]] = 1 + + state._clear_covers() + # Erase all prime markings + state.marked[state.marked == 2] = 0 + return _step3 + + +def _step6(state): + """ + Add the value found in Step 4 to every element of each covered row, + and subtract it from every element of each uncovered column. + Return to Step 4 without altering any stars, primes, or covered lines. + """ + # the smallest uncovered value in the matrix + if np.any(state.row_uncovered) and np.any(state.col_uncovered): + minval = np.min(state.C[state.row_uncovered], axis=0) + minval = np.min(minval[state.col_uncovered]) + state.C[~state.row_uncovered] += minval + state.C[:, state.col_uncovered] -= minval + return _step4 + + +def assignment_problem(costrix): + row_ind, col_ind = linear_sum_assignment(costrix) + return row_ind, col_ind diff --git a/deeplabcut/lmot/main.py b/deeplabcut/lmot/main.py new file mode 100644 index 0000000000..a91095d7af --- /dev/null +++ b/deeplabcut/lmot/main.py @@ -0,0 +1,122 @@ +import os +import os.path +import sys +import time +from pathlib import Path + +import cv2 +import numpy as np +import tensorflow as tf +from skimage.util import img_as_ubyte + +from deeplabcut.lmot import mot +from deeplabcut.lmot.assets import bbox +from deeplabcut.pose_estimation_tensorflow.config import load_config +from deeplabcut.pose_estimation_tensorflow.nnet import predict +from deeplabcut.utils import auxiliaryfunctions + + +def analyze_image(config, vid, shuffle=1, trainingsetindex=0, gputouse=None, save_as_csv=False): + if 'TF_CUDNN_USE_AUTOTUNE' in os.environ: + del os.environ['TF_CUDNN_USE_AUTOTUNE'] # was potentially set during training + + tf.reset_default_graph() + cfg = auxiliaryfunctions.read_config(config) + trainFraction = cfg['TrainingFraction'][trainingsetindex] + + modelfolder = os.path.join(cfg["project_path"], str(auxiliaryfunctions.GetModelFolder(trainFraction, shuffle, cfg))) + path_test_config = Path(modelfolder) / 'test' / 'pose_cfg.yaml' + try: + dlc_cfg = load_config(str(path_test_config)) + except FileNotFoundError: + raise FileNotFoundError( + "It seems the model for shuffle %s and trainFraction %s does not exist." % (shuffle, trainFraction)) + + # Check which snapshots are available and sort them by # iterations + try: + Snapshots = np.array( + [fn.split('.')[0] for fn in os.listdir(os.path.join(modelfolder, 'train')) if "index" in fn]) + except FileNotFoundError: + raise FileNotFoundError( + "Snapshots not found! It seems the dataset for shuffle %s has not been trained/does not exist.\n Please train it before using it to analyze videos.\n Use the function 'train_network' to train the network for shuffle %s." % ( + shuffle, shuffle)) + + if cfg['snapshotindex'] == 'all': + print( + "Snapshotindex is set to 'all' in the config.yaml file. Running video analysis with all snapshots is very costly! Use the function 'evaluate_network' to choose the best the snapshot. For now, changing snapshot index to -1!") + snapshotindex = -1 + else: + snapshotindex = cfg['snapshotindex'] + + increasing_indices = np.argsort([int(m.split('-')[1]) for m in Snapshots]) + Snapshots = Snapshots[increasing_indices] + + print("Using %s" % Snapshots[snapshotindex], "for model", modelfolder) + + ################################################## + # Load and setup CNN part detector + ################################################## + + # Check if data already was generated: + dlc_cfg['init_weights'] = os.path.join(modelfolder, 'train', Snapshots[snapshotindex]) + + # update batchsize (based on parameters in config.yaml) + dlc_cfg['batch_size'] = 1 + + sess, inputs, outputs = predict.setup_pose_prediction(dlc_cfg) + + if gputouse is not None: # gpu selectinon + os.environ['CUDA_VISIBLE_DEVICES'] = str(gputouse) + ##################################################### + # Video analysis + ##################################################### + print("Starting to analyze % ", vid) + ##################################################### + # Read Video + ##################################################### + tracker = mot.Tracker(100) + vname = Path(vid).stem + '.avi' + print("Loading ", vid) + cap = cv2.VideoCapture(vid) + frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + out = cv2.VideoWriter(vname, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 25, (frame_width, frame_height)) + while cap.isOpened(): + ret, frame = cap.read() + start = time.time() + if ret: + scrmap, locref = getPose(dlc_cfg, sess, inputs, outputs, frame) + tracker.track(scrmap, locref) + for obj in range(len(tracker.tracks)): + bbox(frame, tracker.tracks[obj]) + out.write(frame) + else: + print("frame was analyzed in " + str(time.time() - start)) + print("") + print("") + break + # cv2.waitKey(50) + # do a bit of cleanup + print("[INFO] cleaning up...") + cv2.destroyAllWindows() + out.release() + + +def getPose(dlc_cfg, sess, inputs, outputs, img): + frame = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + frame = img_as_ubyte(frame) + scmap, locref = predict.getpose(frame, dlc_cfg, sess, inputs, outputs, live=True) + return scmap, locref + + +if __name__ == '__main__': + if len(sys.argv) == 3: + conf = sys.argv[1] + vid = sys.argv[2] + print('configuration :', conf) + print('input vid :', vid) + analyze_image(conf, vid) + else: + print('Usage: python3 main.py [config path] [image]\n' + 'config path: configuration path of the trained network\n' + 'video: input video') diff --git a/deeplabcut/lmot/mot.py b/deeplabcut/lmot/mot.py new file mode 100644 index 0000000000..f2312185c5 --- /dev/null +++ b/deeplabcut/lmot/mot.py @@ -0,0 +1,82 @@ +import numpy as np + +from deeplabcut.lmot.extract_local_maxima import extract_locmaxima +from deeplabcut.lmot.linear_sum_assignment import assignment_problem +from deeplabcut.lmot.track import Track + + +class Tracker: + def __init__(self, max_frames_to_skip): + self.max_frames_to_skip = max_frames_to_skip + self.id_count = 1 + self.tracks = [] + self.iter = 0 + self.det = [] + + 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)): + gating_dist = self.tracks[i].KF.mahalanobis_dist(detections[j].coordinates) + cost_matrix[i][j] = gating_dist + return cost_matrix + + def track(self, scmap, locref): + print("ITERATION NUMBER : " + str(self.iter)) + self.iter += 1 + scmap = np.squeeze(scmap) + detections = extract_locmaxima(scmap, locref) + self.det = detections + + if len(self.tracks) == 0: + # print("# No Tracks Found, Creating New Tracks") + for i in range(len(detections)): + track = Track(detections[i], self.id_count) + self.id_count += 1 + self.tracks.append(track) + # print("# number of created tracks: " + str(len(self.tracks))) + N = len(self.tracks) + M = len(detections) + print(detections, N, M) + cost_matrix = self.cost_metric(detections, N, M) + rows, cols = assignment_problem(cost_matrix) + matches = self._matching(rows, cols, cost_matrix, detections, N) + for i in range(len(matches)): + if matches[i] is not None: + self.tracks[i].skipped_frames = 0 + self.tracks[i].predict(self.iter, detections[matches[i]]) + else: + self.tracks[i].predict(it=self.iter) + print("tracks at this iteration", self.tracks, "counting", len(self.tracks)) + + def _matching(self, rows, cols, costmatrix, detections, N): + c = 0 + matches = [None] * N + for i in range(len(rows)): + matches[rows[i]] = cols[i] + c += 1 + unmatching = [] + for i in range(len(matches)): + if matches[i] is not None: + if costmatrix[i][matches[i]] > 150: + print("unmatched found") + matches[i] = None + unmatching.append(i) + else: + self.tracks[i].skipped_frames += 1 + del_tracks = [] + for i in range(len(self.tracks)): + if self.tracks[i].skipped_frames > self.max_frames_to_skip: + del_tracks.append(self.tracks[i]) + + if len(del_tracks) > 0: + for i in range(len(del_tracks)): + del self.tracks[i] + del matches[i] + + for i in range(len(detections)): + if i not in matches: + track = Track(detections[i], self.id_count) + self.id_count += 1 + self.tracks.append(track) + return matches diff --git a/deeplabcut/lmot/notes.txt b/deeplabcut/lmot/notes.txt new file mode 100644 index 0000000000..9ce80212f9 --- /dev/null +++ b/deeplabcut/lmot/notes.txt @@ -0,0 +1,10 @@ +Note: This was done in the course of my studies as a Bachelor's project in the Straw Lab at the University of Freiburg. +The Straw Lab(https://strawlab.org/) is interested in real-time flies tracking, so the cost metric +& other stuff were specifically chosen with the regard to the input. + +More Notes +- The Kalman filter is a constant velocity model + +- Evaluation was done using only the initial input videos (See picture in header) + +- This feature merely tested on DLC2, as it was developed under DLC1 \ No newline at end of file diff --git a/deeplabcut/lmot/track.py b/deeplabcut/lmot/track.py new file mode 100644 index 0000000000..c16efb10e6 --- /dev/null +++ b/deeplabcut/lmot/track.py @@ -0,0 +1,21 @@ +from deeplabcut.lmot.assets import get_random_color +from deeplabcut.lmot.kalman_filter import KalmanFilter + + +class Track: + def __init__(self, detection, trackId): + self.track_id = trackId + self.KF = KalmanFilter(detection.coordinates) + self.prediction = detection.coordinates + self.height = detection.height + self.width = detection.width + self.trace = dict() # trace path + self.skipped_frames = 0 + self.color = get_random_color(self.track_id) + + def predict(self, it, detection=None): + self.KF.predict() + if detection is not None: + self.KF.update(detection.coordinates) + self.prediction = detection.coordinates + self.trace[it] = self.prediction diff --git a/deeplabcut/pose_estimation_tensorflow/nnet/predict.py b/deeplabcut/pose_estimation_tensorflow/nnet/predict.py index 4b20d2e3f0..7424cc11f7 100644 --- a/deeplabcut/pose_estimation_tensorflow/nnet/predict.py +++ b/deeplabcut/pose_estimation_tensorflow/nnet/predict.py @@ -18,6 +18,7 @@ import numpy as np import tensorflow as tf + vers = (tf.__version__).split('.') if int(vers[0])==1 and int(vers[1])>12: TF=tf.compat.v1 @@ -25,6 +26,7 @@ TF=tf from deeplabcut.pose_estimation_tensorflow.nnet.net_factory import pose_net + def setup_pose_prediction(cfg): TF.reset_default_graph() inputs = TF.placeholder(tf.float32, shape=[cfg.batch_size , None, None, 3]) @@ -53,10 +55,11 @@ def extract_cnn_output(outputs_np, cfg): shape = locref.shape locref = np.reshape(locref, (shape[0], shape[1], -1, 2)) locref *= cfg.locref_stdev - if len(scmap.shape)==2: #for single body part! - scmap=np.expand_dims(scmap,axis=2) + if len(scmap.shape) == 2: # for single body part! or live mode ? + scmap = np.expand_dims(scmap, axis=2) return scmap, locref + def argmax_pose_predict(scmap, offmat, stride): """Combine scoremat and offsets to the final pose.""" num_joints = scmap.shape[2] @@ -71,6 +74,23 @@ def argmax_pose_predict(scmap, offmat, stride): [scmap[maxloc][joint_idx]]))) return np.array(pose) + +# modified by: Tabet Ehsainieh +def getpose(image, cfg, sess, inputs, outputs, outall=False, live=False): + ''' Extract pose ''' + im = np.expand_dims(image, axis=0).astype(float) + outputs_np = sess.run(outputs, feed_dict={inputs: im}) + scmap, locref = extract_cnn_output(outputs_np, cfg) + # real time tracking mode + if live: + return scmap, locref + # if scmap.shape[2] == 1: + # return scmap, locref + # else: + # raise Exception("looks like your chosen configuration meant for different purpose " + # "... existing live tracking mode") + pose = argmax_pose_predict(scmap, locref, cfg.stride) + def multi_pose_predict(scmap, locref, stride, num_outputs): Y, X = get_top_values(scmap[None], num_outputs) Y, X = Y[:, 0], X[:, 0] @@ -109,19 +129,20 @@ def getpose(image, cfg, sess, inputs, outputs, outall=False): else: return pose -## Functions below implement are for batch sizes > 1: + +# Functions below implement are for batch sizes > 1 def extract_cnn_outputmulti(outputs_np, cfg): ''' extract locref + scmap from network Dimensions: image batch x imagedim1 x imagedim2 x bodypart''' scmap = outputs_np[0] locref = None if cfg.location_refinement: - locref =outputs_np[1] + locref = outputs_np[1] shape = locref.shape - locref = np.reshape(locref, (shape[0], shape[1],shape[2], -1, 2)) + locref = np.reshape(locref, (shape[0], shape[1], shape[2], -1, 2)) locref *= cfg.locref_stdev - if len(scmap.shape)==2: #for single body part! - scmap=np.expand_dims(scmap,axis=2) + if len(scmap.shape) == 2: # for single body part! + scmap = np.expand_dims(scmap, axis=2) return scmap, locref @@ -148,6 +169,25 @@ def getposeNP(image, cfg, sess, inputs, outputs, outall=False): num_outputs = cfg.get('num_outputs', 1) outputs_np = sess.run(outputs, feed_dict={inputs: image}) + scmap, locref = extract_cnn_outputmulti(outputs_np, cfg) # processes image batch. + batchsize, ny, nx, num_joints = scmap.shape + + # Combine scoremat and offsets to the final pose. + LOCREF = locref.reshape(batchsize, nx * ny, num_joints, 2) + MAXLOC = np.argmax(scmap.reshape(batchsize, nx * ny, num_joints), axis=1) + Y, X = np.unravel_index(MAXLOC, dims=(ny, nx)) + DZ = np.zeros((batchsize, num_joints, 3)) + for l in range(batchsize): + for k in range(num_joints): + DZ[l, k, :2] = LOCREF[l, MAXLOC[l, k], k, :] + DZ[l, k, 2] = scmap[l, Y[l, k], X[l, k], k] + + X = X.astype('float32') * cfg.stride + .5 * cfg.stride + DZ[:, :, 0] + Y = Y.astype('float32') * cfg.stride + .5 * cfg.stride + DZ[:, :, 1] + pose = np.empty((cfg['batch_size'], cfg['num_joints'] * 3), dtype=X.dtype) + pose[:, 0::3] = X + pose[:, 1::3] = Y + pose[:, 2::3] = DZ[:, :, 2] # P scmap, locref = extract_cnn_outputmulti(outputs_np, cfg) #processes image batch. batchsize,ny,nx,num_joints = scmap.shape @@ -200,4 +240,4 @@ def setup_GPUpose_prediction(cfg): return sess, inputs, outputs def extract_GPUprediction(outputs, cfg): - return outputs[0] + return outputs[0] \ No newline at end of file