DeepLabCut · ehsainit · Dec 21, 2019 · Dec 21, 2019 · Dec 21, 2019 · Dec 21, 2019
diff --git a/deeplabcut/lmot/__init__.py b/deeplabcut/lmot/__init__.py
diff --git a/deeplabcut/lmot/assets.py 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
@@ -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
@@ -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
@@ -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