track.py

import cv2
import numpy as np
import random
from scipy.spatial import distance
from tensorflow import keras

import settings


class Track:
    """ Class containing information about a single track. """

    def __init__(self, center_x : float, center_y : float,
                 width : float, height : float) -> None:
        self.vector = np.array([center_x, center_y, width / height, height, 0, 0, 0, 0])
        self.age = 0
        self.identifier = random.randint(1000, 9999)  # TODO: Make it unique
        self._vector = np.array([center_x, center_y, width / height, height, 0, 0, 0, 0])
        self.movement_model = KalmanFilter(
            x=self._vector, P=np.eye(8),
            F=np.array([
                [1, 0, 0, 0, settings.TIME_STEP, 0, 0, 0],
                [0, 1, 0, 0, 0, settings.TIME_STEP, 0, 0],
                [0, 0, 1, 0, 0, 0, settings.TIME_STEP, 0],
                [0, 0, 0, 1, 0, 0, 0, settings.TIME_STEP],
                [0, 0, 0, 0, 1, 0, 0, 0],
                [0, 0, 0, 0, 0, 1, 0, 0],
                [0, 0, 0, 0, 0, 0, 1, 0],
                [0, 0, 0, 0, 0, 0, 0, 1],
            ]),
            B=np.zeros((8, 8)),  # control input model is zero
            H=np.eye(4, 8),  # we map 8x1 to 4x1
            Q = np.diag(settings.KALMAN_NOISE_COVARIANCE_Q),
            R = np.diag(settings.KALMAN_NOISE_COVARIANCE_R)
        )
        self.descriptors = []
        self.update_count = 0

    @property
    def rect(self) -> tuple:
        x = int(self.vector[0] - (self.vector[2] * self.vector[3]) // 2)
        y = int(self.vector[1] - self.vector[3] // 2)
        return (x if x >= 0 else 0, y if y >= 0 else 0, int(self.vector[2] * self.vector[3]), self.vector[3])

    @property
    def topleft(self) -> tuple:
        return (int(self.rect[0]), int(self.rect[1]))

    @property
    def bottomright(self) -> tuple:
        return (int(self.rect[0] + int(self.vector[2] * self.vector[3])), int(self.rect[1] + self.vector[3]))

    def mahalanobis_distance(self, detection) -> float:
        covariance_matrix = np.cov(np.array([self._vector[:4], detection.vector[:4]]).T) + 0.00001 * np.eye(4)
        inverted_S = np.linalg.inv(covariance_matrix)
        diff = self._vector[:4] - detection.vector[:4]
        return np.sqrt((diff.T @ inverted_S) @ diff)

    def threshold_mahalanobis_distance(self, detection) -> float:
        value = self.mahalanobis_distance(detection)
        return 1 if value <= settings.MAHALANOBIS_THRESHOLD else 0

    def smallest_cosine_distance(self, detection) -> float:
        return min([1 - detection.descriptors[-1].T @ descriptor for descriptor in self.descriptors])

    def threshold_smallest_cosine_distance(self, detection) -> float:
        value = self.smallest_cosine_distance(detection)
        return 1 if value <= settings.SMALLEST_COSINE_THRESHOLD else 0

    def calculate_cost(self, detection, weight) -> float:
        return weight * self.mahalanobis_distance(detection) + \
            (1 - weight) * self.smallest_cosine_distance(detection)

    def calculate_gate(self, detection) -> float:
        return self.threshold_mahalanobis_distance(detection) * \
            self.threshold_smallest_cosine_distance(detection)

    def update_descriptors(self, frame, resnet) -> None:
        capped_x_max = self.bottomright[0] if self.bottomright[0] < frame.shape[1] else frame.shape[1] - 1
        capped_y_max = self.bottomright[1] if self.bottomright[1] < frame.shape[0] else frame.shape[0] - 1
        image = frame[self.topleft[1]:capped_y_max, self.topleft[0]:capped_x_max]
        image = cv2.resize(image, settings.RESNET_RESIZE_SIZE)
        image = np.expand_dims(image, axis=0)
        image = keras.applications.resnet50.preprocess_input(image)
        features = resnet.predict(image, verbose=False)
        features = features.flatten()
        features_magnitude = np.linalg.norm(features)  # Magnitude of the vector for normalizing it
        features = features / features_magnitude
        if len(self.descriptors) < settings.MAX_DESCRIPTORS:
            self.descriptors.append(features)
        else:
            self.descriptors = self.descriptors[1:]
            self.descriptors.append(features)

    def update(self, detection, frame, resnet) -> None:
        if detection:
            self.vector = detection.vector
            self.age = 0
            self._vector, P = self.movement_model.update(detection.vector[:4])
            if len(self.descriptors) < settings.MAX_DESCRIPTORS:
                self.descriptors.append(detection.descriptors[-1])
            else:
                self.descriptors = self.descriptors[1:]
                self.descriptors.append(detection.descriptors[-1])
        else:
            self.age += 1
        self._vector, P = self.movement_model.predict(self._vector)
        self.update_count += 1

    def draw_rect(self, image, color) -> None:
        cv2.rectangle(image, self.topleft, self.bottomright, color, 1)
        cv2.putText(image, str(self.identifier), (self.rect[0] - 10, self.rect[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)


class KalmanFilter:
    """ Class for continuous prediction using Kalman Filter. The movement model. """

    def __init__(self, x : np.ndarray, P : np.ndarray, F : np.ndarray,
                 B : np.ndarray, H : np.ndarray, Q : np.ndarray,
                 R : np.ndarray) -> None:
        self.x = x  # initial state
        self.P = P  # covariance of initial state
        self.F = F  # the model
        self.B = B  # input control
        self.H = H  # mapping between observation space and prediction space
        self.Q = Q  # covariance noise
        self.R = R  # covariance noise

    def predict(self, u : np.ndarray) -> np.ndarray:
        self.x = self.F @ self.x + self.B @ u
        self.P = (self.F @ self.P) @ self.F.T + self.Q
        return self.x, self.P

    def update(self, z : np.ndarray) -> np.ndarray:
        S = (self.H @ self.P) @ self.H.T + self.R
        K = (self.P @ self.H.T) @ np.linalg.inv(S)  # inverse of S
        y = z - (self.H @ self.x)
        self.x = self.x + (K @ y)
        I = np.eye(self.P.shape[0])
        self.P = ((I - (K @ self.H)) @ self.P) @ ((I - (K @ self.H).T) + (K @ self.R) @ K.T)
        return self.x, self.P