From fb8f2868090c292ac9c24c5d9426fcf9253173b1 Mon Sep 17 00:00:00 2001
From: ehsainit <tabet.ehsainieh@gmx.de>
Date: Sat, 21 Dec 2019 14:15:24 +0100
Subject: [PATCH 1/8] new feature: adding real-time multiple object tracking
 DLC-inference-based

---
 deeplabcut/lmot/__init__.py                   |   0
 deeplabcut/lmot/assets.py                     |  30 ++
 deeplabcut/lmot/detection.py                  |   9 +
 deeplabcut/lmot/extract_local_maxima.py       |  40 +++
 deeplabcut/lmot/kalman_filter.py              | 108 +++++++
 deeplabcut/lmot/linear_sum_assignment.py      | 287 ++++++++++++++++++
 deeplabcut/lmot/main.py                       | 126 ++++++++
 deeplabcut/lmot/mot.py                        |  87 ++++++
 deeplabcut/lmot/track.py                      |  25 ++
 .../nnet/predict.py                           |  76 +++--
 10 files changed, 757 insertions(+), 31 deletions(-)
 create mode 100644 deeplabcut/lmot/__init__.py
 create mode 100644 deeplabcut/lmot/assets.py
 create mode 100644 deeplabcut/lmot/detection.py
 create mode 100644 deeplabcut/lmot/extract_local_maxima.py
 create mode 100644 deeplabcut/lmot/kalman_filter.py
 create mode 100644 deeplabcut/lmot/linear_sum_assignment.py
 create mode 100644 deeplabcut/lmot/main.py
 create mode 100644 deeplabcut/lmot/mot.py
 create mode 100644 deeplabcut/lmot/track.py

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..916a01ada2
--- /dev/null
+++ b/deeplabcut/lmot/assets.py
@@ -0,0 +1,30 @@
+# Copyright 2019 by
+# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
+# All rights reserved.
+
+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..c34ddedc5d
--- /dev/null
+++ b/deeplabcut/lmot/detection.py
@@ -0,0 +1,9 @@
+# Copyright 2019 by
+# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
+# All rights reserved.
+
+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..f0a690a997
--- /dev/null
+++ b/deeplabcut/lmot/extract_local_maxima.py
@@ -0,0 +1,40 @@
+# Copyright 2019 by
+# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
+# All rights reserved.
+
+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..bb31ad99c0
--- /dev/null
+++ b/deeplabcut/lmot/kalman_filter.py
@@ -0,0 +1,108 @@
+# Copyright 2019 by Tabet Ehsainieh.
+# All rights reserved.
+# Please see the LICENSE file that should have been included as part of this package.
+# Copyright 2019 by
+# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
+# All rights reserved.
+
+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 <bmc@clapper.org>, 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..1ca418afd4
--- /dev/null
+++ b/deeplabcut/lmot/main.py
@@ -0,0 +1,126 @@
+# Copyright 2019 by
+# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
+# All rights reserved.
+
+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..acd67b3444
--- /dev/null
+++ b/deeplabcut/lmot/mot.py
@@ -0,0 +1,87 @@
+# Copyright 2019 by
+# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
+# All rights reserved.
+
+
+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/track.py b/deeplabcut/lmot/track.py
new file mode 100644
index 0000000000..2458e5e59c
--- /dev/null
+++ b/deeplabcut/lmot/track.py
@@ -0,0 +1,25 @@
+# Copyright 2019 by
+# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
+# All rights reserved.
+
+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 2fd15def18..8904294764 100644
--- a/deeplabcut/pose_estimation_tensorflow/nnet/predict.py
+++ b/deeplabcut/pose_estimation_tensorflow/nnet/predict.py
@@ -8,11 +8,13 @@
 
 import numpy as np
 import tensorflow as 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])
+    inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
     net_heads = pose_net(cfg).test(inputs)
     outputs = [net_heads['part_prob']]
     if cfg.location_refinement:
@@ -27,7 +29,8 @@ def setup_pose_prediction(cfg):
     restorer.restore(sess, cfg.init_weights)
 
     return sess, inputs, outputs
-    
+
+
 def extract_cnn_output(outputs_np, cfg):
     ''' extract locref + scmap from network '''
     scmap = outputs_np[0]
@@ -38,10 +41,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]
@@ -56,58 +60,68 @@ def argmax_pose_predict(scmap, offmat, stride):
                                [scmap[maxloc][joint_idx]])))
     return np.array(pose)
 
-def getpose(image, cfg, sess, inputs, outputs, outall=False):
-    ''' Extract pose '''
-    im=np.expand_dims(image, axis=0).astype(float)
+
+# 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)
     if outall:
         return scmap, locref, pose
     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
 
 
 def getposeNP(image, cfg, sess, inputs, outputs, outall=False):
     ''' Adapted from DeeperCut, performs numpy-based faster inference on batches'''
     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))
+
+    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
+            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
     if outall:
         return scmap, locref, pose
     else:
         return pose
-

From a0cb92df3f50db7239d5a2c3fa20ed1f4b653f50 Mon Sep 17 00:00:00 2001
From: ehsainit <tabet.ehsainieh@gmx.de>
Date: Sat, 21 Dec 2019 14:16:36 +0100
Subject: [PATCH 2/8] new feature: adding real-time multiple object tracking
 DLC-inference-based

---
 deeplabcut/pose_estimation_tensorflow/nnet/predict.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeplabcut/pose_estimation_tensorflow/nnet/predict.py b/deeplabcut/pose_estimation_tensorflow/nnet/predict.py
index 8904294764..da884d6523 100644
--- a/deeplabcut/pose_estimation_tensorflow/nnet/predict.py
+++ b/deeplabcut/pose_estimation_tensorflow/nnet/predict.py
@@ -61,7 +61,7 @@ def argmax_pose_predict(scmap, offmat, stride):
     return np.array(pose)
 
 
-# modified by : Tabet Ehsainieh
+# 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)

From 704c46380a39d928becf3b396a283832f720cc66 Mon Sep 17 00:00:00 2001
From: ehsainit <45032503+ehsainit@users.noreply.github.com>
Date: Sat, 21 Dec 2019 14:18:23 +0100
Subject: [PATCH 3/8] updating copy rights

---
 deeplabcut/lmot/kalman_filter.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/deeplabcut/lmot/kalman_filter.py b/deeplabcut/lmot/kalman_filter.py
index bb31ad99c0..7d19de6a1b 100644
--- a/deeplabcut/lmot/kalman_filter.py
+++ b/deeplabcut/lmot/kalman_filter.py
@@ -1,6 +1,3 @@
-# Copyright 2019 by Tabet Ehsainieh.
-# All rights reserved.
-# Please see the LICENSE file that should have been included as part of this package.
 # Copyright 2019 by
 # Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
 # All rights reserved.

From 3a587a72a562b700b8fd6294cdf6649b3c79b11a Mon Sep 17 00:00:00 2001
From: ehsainit <45032503+ehsainit@users.noreply.github.com>
Date: Thu, 9 Jan 2020 12:06:14 +0100
Subject: [PATCH 4/8] Update README.md

---
 README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/README.md b/README.md
index e0d639808c..65b7779a76 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,7 @@ For a step-by-step user guide, please read the [Nature Protocols paper](https://
 # [DEMO the code](/examples)
 
 We provide several Jupyter Notebooks: one that walks you through a demo dataset to test your installation, and another Notebook to run DeepLabCut from the beginning on your own data. We also show you how to use the code in Docker, and on Google Colab.
+# Update: DeepLabCut with new feature: Real-Time MOT tracking 
 
 # Why use DeepLabCut?
 

From dd3d66996edf3835970d8732dfad06ec2f58c607 Mon Sep 17 00:00:00 2001
From: ehsainit <tabet.ehsainieh@gmx.de>
Date: Thu, 30 Jan 2020 23:37:01 +0100
Subject: [PATCH 5/8] removing headers

---
 deeplabcut/lmot/assets.py               | 4 ----
 deeplabcut/lmot/detection.py            | 4 ----
 deeplabcut/lmot/extract_local_maxima.py | 4 ----
 deeplabcut/lmot/kalman_filter.py        | 4 ----
 deeplabcut/lmot/main.py                 | 4 ----
 deeplabcut/lmot/mot.py                  | 5 -----
 deeplabcut/lmot/track.py                | 4 ----
 7 files changed, 29 deletions(-)

diff --git a/deeplabcut/lmot/assets.py b/deeplabcut/lmot/assets.py
index 916a01ada2..3d3784950b 100644
--- a/deeplabcut/lmot/assets.py
+++ b/deeplabcut/lmot/assets.py
@@ -1,7 +1,3 @@
-# Copyright 2019 by
-# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
-# All rights reserved.
-
 import random
 
 import cv2
diff --git a/deeplabcut/lmot/detection.py b/deeplabcut/lmot/detection.py
index c34ddedc5d..8a9fb4d7e0 100644
--- a/deeplabcut/lmot/detection.py
+++ b/deeplabcut/lmot/detection.py
@@ -1,7 +1,3 @@
-# Copyright 2019 by
-# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
-# All rights reserved.
-
 class Detection:
     def __init__(self, coordinates, height, width):
         self.coordinates = coordinates
diff --git a/deeplabcut/lmot/extract_local_maxima.py b/deeplabcut/lmot/extract_local_maxima.py
index f0a690a997..0fa3bf72a3 100644
--- a/deeplabcut/lmot/extract_local_maxima.py
+++ b/deeplabcut/lmot/extract_local_maxima.py
@@ -1,7 +1,3 @@
-# Copyright 2019 by
-# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
-# All rights reserved.
-
 import numpy as np
 import scipy.ndimage as ndimage
 import scipy.ndimage.filters as filters
diff --git a/deeplabcut/lmot/kalman_filter.py b/deeplabcut/lmot/kalman_filter.py
index 7d19de6a1b..b3f84f65eb 100644
--- a/deeplabcut/lmot/kalman_filter.py
+++ b/deeplabcut/lmot/kalman_filter.py
@@ -1,7 +1,3 @@
-# Copyright 2019 by
-# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
-# All rights reserved.
-
 from copy import deepcopy
 
 import numpy as np
diff --git a/deeplabcut/lmot/main.py b/deeplabcut/lmot/main.py
index 1ca418afd4..a91095d7af 100644
--- a/deeplabcut/lmot/main.py
+++ b/deeplabcut/lmot/main.py
@@ -1,7 +1,3 @@
-# Copyright 2019 by
-# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
-# All rights reserved.
-
 import os
 import os.path
 import sys
diff --git a/deeplabcut/lmot/mot.py b/deeplabcut/lmot/mot.py
index acd67b3444..f2312185c5 100644
--- a/deeplabcut/lmot/mot.py
+++ b/deeplabcut/lmot/mot.py
@@ -1,8 +1,3 @@
-# Copyright 2019 by
-# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
-# All rights reserved.
-
-
 import numpy as np
 
 from deeplabcut.lmot.extract_local_maxima import extract_locmaxima
diff --git a/deeplabcut/lmot/track.py b/deeplabcut/lmot/track.py
index 2458e5e59c..c16efb10e6 100644
--- a/deeplabcut/lmot/track.py
+++ b/deeplabcut/lmot/track.py
@@ -1,7 +1,3 @@
-# Copyright 2019 by
-# Tabet Ehsainieh, ehsainit@informatik.uni-freiburg.de
-# All rights reserved.
-
 from deeplabcut.lmot.assets import get_random_color
 from deeplabcut.lmot.kalman_filter import KalmanFilter
 

From 1dc088f1e5ab4164fc646d5f1129437a675038bc Mon Sep 17 00:00:00 2001
From: ehsainit <tabet.ehsainieh@gmx.de>
Date: Thu, 30 Jan 2020 23:56:54 +0100
Subject: [PATCH 6/8] README.md remove line

---
 README.md | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 65b7779a76..c120a41c86 100644
--- a/README.md
+++ b/README.md
@@ -46,8 +46,7 @@ For a step-by-step user guide, please read the [Nature Protocols paper](https://
 
 # [DEMO the code](/examples)
 
-We provide several Jupyter Notebooks: one that walks you through a demo dataset to test your installation, and another Notebook to run DeepLabCut from the beginning on your own data. We also show you how to use the code in Docker, and on Google Colab.
-# Update: DeepLabCut with new feature: Real-Time MOT tracking 
+We provide several Jupyter Notebooks: one that walks you through a demo dataset to test your installation, and another Notebook to run DeepLabCut from the beginning on your own data. We also show you how to use the code in Docker, and on Google Colab. 
 
 # Why use DeepLabCut?
 

From c34b849999a497f81675f9c27ef4d212e71c7183 Mon Sep 17 00:00:00 2001
From: ehsainit <tabet.ehsainieh@gmx.de>
Date: Mon, 16 Mar 2020 21:32:23 +0100
Subject: [PATCH 7/8] general notes and information

---
 deeplabcut/lmot/notes.txt | 10 ++++++++++
 1 file changed, 10 insertions(+)
 create mode 100644 deeplabcut/lmot/notes.txt

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

From fbdede4261d7924592d5813a4e97e9a4f85c9ca4 Mon Sep 17 00:00:00 2001
From: Mackenzie Mathis <mathis@rowland.harvard.edu>
Date: Sat, 21 Mar 2020 18:27:51 -0400
Subject: [PATCH 8/8] fix space

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index c120a41c86..e0d639808c 100644
--- a/README.md
+++ b/README.md
@@ -46,7 +46,7 @@ For a step-by-step user guide, please read the [Nature Protocols paper](https://
 
 # [DEMO the code](/examples)
 
-We provide several Jupyter Notebooks: one that walks you through a demo dataset to test your installation, and another Notebook to run DeepLabCut from the beginning on your own data. We also show you how to use the code in Docker, and on Google Colab. 
+We provide several Jupyter Notebooks: one that walks you through a demo dataset to test your installation, and another Notebook to run DeepLabCut from the beginning on your own data. We also show you how to use the code in Docker, and on Google Colab.
 
 # Why use DeepLabCut?