deep_ocsort.py

"""
    https://github.com/mikel-brostrom/yolo_tracking.git
    This script is adopted from the SORT script by Alex Bewley alex@bewley.ai
    Additional changes from https://github.com/mikel-brostrom/yolo_tracking:
        - Removed unused functions
        - Modified ReIDDetectMultiBackend to use YoutuReid (OpenCV Zoo)
        - Use segmented person on black image instead of box crop for embeddings
        - Included behavior analysis (emotion, gender, posture)
        - KalmaxBoxTracker additional parameters:
            - Added raw person embeddings for later multicamera association
            - Emotion, gender, posture
        - Third round of OCSORT association by comparing raw feature embedding 
            (might be useful for reappearance but yet to be fully investigated)
            - needs hyperparameter tuning (reid_thresh)
"""
import sys

import cv2
import torch
import torchreid
import numpy as np

from association import *
from kalman_filter import KalmanFilterNew
from cmc import CameraMotionCompensation


def k_previous_obs(observations, cur_age, k):
    if len(observations) == 0:
        return [-1, -1, -1, -1, -1]
    for i in range(k):
        dt = k - i
        if cur_age - dt in observations:
            return observations[cur_age - dt]
    max_age = max(observations.keys())
    return observations[max_age]


def convert_bbox_to_z_new(bbox):
    w = bbox[2] - bbox[0]
    h = bbox[3] - bbox[1]
    x = bbox[0] + w / 2.0
    y = bbox[1] + h / 2.0
    return np.array([x, y, w, h]).reshape((4, 1))


def convert_x_to_bbox_new(x):
    x, y, w, h = x.reshape(-1)[:4]
    return np.array([x - w / 2, y - h / 2, x + w / 2, y + h / 2]).reshape(1, 4)


def speed_direction(bbox1, bbox2):
    cx1, cy1 = (bbox1[0] + bbox1[2]) / 2.0, (bbox1[1] + bbox1[3]) / 2.0
    cx2, cy2 = (bbox2[0] + bbox2[2]) / 2.0, (bbox2[1] + bbox2[3]) / 2.0
    speed = np.array([cy2 - cy1, cx2 - cx1])
    norm = np.sqrt((cy2 - cy1) ** 2 + (cx2 - cx1) ** 2) + 1e-6
    return speed / norm


def new_kf_process_noise(w, h, p=1 / 20, v=1 / 160):
    Q = np.diag(
        (
            (p * w) ** 2,
            (p * h) ** 2,
            (p * w) ** 2,
            (p * h) ** 2,
            (v * w) ** 2,
            (v * h) ** 2,
            (v * w) ** 2,
            (v * h) ** 2,
        )
    )
    return Q


def new_kf_measurement_noise(w, h, m=1 / 20):
    w_var = (m * w) ** 2
    h_var = (m * h) ** 2
    R = np.diag((w_var, h_var, w_var, h_var))
    return R


class KalmanBoxTracker(object):
    """
    This class represents the internal state of individual tracked objects observed as bbox.
    """

    count = 0

    def __init__(self, bbox, cls, behavior, delta_t=3, emb=None, alpha=0):
        """
        Initialises a tracker using initial bounding box.

        """
        # define constant velocity model
        self.cls = cls
        self.conf = bbox[-1]

        self.kf = KalmanFilterNew(dim_x=8, dim_z=4)
        self.kf.F = np.array(
            [
                # x y w h x' y' w' h'
                [1, 0, 0, 0, 1, 0, 0, 0],
                [0, 1, 0, 0, 0, 1, 0, 0],
                [0, 0, 1, 0, 0, 0, 1, 0],
                [0, 0, 0, 1, 0, 0, 0, 1],
                [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],
            ]
        )
        self.kf.H = np.array(
            [
                [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, 0, 0, 0, 0],
            ]
        )
        _, _, w, h = convert_bbox_to_z_new(bbox).reshape(-1)
        self.kf.P = new_kf_process_noise(w, h)
        self.kf.P[:4, :4] *= 4
        self.kf.P[4:, 4:] *= 100
        # Process and measurement uncertainty happen in functions
        self.bbox_to_z_func = convert_bbox_to_z_new
        self.x_to_bbox_func = convert_x_to_bbox_new

        self.kf.x[:4] = self.bbox_to_z_func(bbox)

        self.time_since_update = 0
        self.id = KalmanBoxTracker.count
        KalmanBoxTracker.count += 1
        self.history = []
        self.hits = 0
        self.hit_streak = 0
        self.age = 0
        """
        NOTE: [-1,-1,-1,-1,-1] is a compromising placeholder for non-observation status, the same for the return of 
        function k_previous_obs. It is ugly and I do not like it. But to support generate observation array in a 
        fast and unified way, which you would see below k_observations = np.array([k_previous_obs(...]]), let's bear it for now.
        """
        # Used for OCR
        self.last_observation = np.array([-1, -1, -1, -1, -1])  # placeholder
        # Used to output track after min_hits reached
        self.history_observations = []
        # Used for velocity
        self.observations = dict()
        self.velocity = None
        self.delta_t = delta_t

        self.emb = emb
        self.emb_raw = emb
        self.highest_conf = bbox[-1]
        self.behavior = behavior

        self.frozen = False

    def update(self, bbox, cls):
        """
        Updates the state vector with observed bbox.
        """
        if bbox is not None:
            self.frozen = False
            self.cls = cls
            self.conf = bbox[-1]
            if self.last_observation.sum() >= 0:  # no previous observation
                previous_box = None
                for dt in range(self.delta_t, 0, -1):
                    if self.age - dt in self.observations:
                        previous_box = self.observations[self.age - dt]
                        break
                if previous_box is None:
                    previous_box = self.last_observation
                """
                  Estimate the track speed direction with observations \Delta t steps away
                """
                self.velocity = speed_direction(previous_box, bbox)
            """
              Insert new observations. This is a ugly way to maintain both self.observations
              and self.history_observations. Bear it for the moment.
            """
            self.last_observation = bbox
            self.observations[self.age] = bbox
            self.history_observations.append(bbox)

            self.time_since_update = 0
            self.history = []
            self.hits += 1
            self.hit_streak += 1
            R = new_kf_measurement_noise(self.kf.x[2, 0], self.kf.x[3, 0])
            self.kf.update(self.bbox_to_z_func(bbox), R=R)
        else:
            self.kf.update(bbox)
            self.frozen = True

    def update_emb(self, emb, alpha=0.9):
        self.emb = alpha * self.emb + (1 - alpha) * emb
        self.emb /= np.linalg.norm(self.emb)

    def update_emb_raw(self, emb, score):
        if score > self.highest_conf:
            self.emb_raw = emb
            self.highest_conf = score

    def update_behavior(self, behavior):
        self.behavior = behavior

    def get_emb(self):
        return self.emb.cpu()

    def get_emb_raw(self):
        return self.emb_raw.cpu()

    def apply_affine_correction(self, affine):
        m = affine[:, :2]
        t = affine[:, 2].reshape(2, 1)
        # For OCR
        if self.last_observation.sum() > 0:
            ps = self.last_observation[:4].reshape(2, 2).T
            ps = m @ ps + t
            self.last_observation[:4] = ps.T.reshape(-1)

        # Apply to each box in the range of velocity computation
        for dt in range(self.delta_t, -1, -1):
            if self.age - dt in self.observations:
                ps = self.observations[self.age - dt][:4].reshape(2, 2).T
                ps = m @ ps + t
                self.observations[self.age - dt][:4] = ps.T.reshape(-1)

        # Also need to change kf state, but might be frozen
        self.kf.apply_affine_correction(m, t, True)

    def predict(self):
        """
        Advances the state vector and returns the predicted bounding box estimate.
        """
        # Don't allow negative bounding boxes
        if self.kf.x[2] + self.kf.x[6] <= 0:
            self.kf.x[6] = 0
        if self.kf.x[3] + self.kf.x[7] <= 0:
            self.kf.x[7] = 0

        # Stop velocity, will update in kf during OOS
        if self.frozen:
            self.kf.x[6] = self.kf.x[7] = 0
        Q = new_kf_process_noise(self.kf.x[2, 0], self.kf.x[3, 0])

        self.kf.predict(Q=Q)
        self.age += 1
        if self.time_since_update > 0:
            self.hit_streak = 0
        self.time_since_update += 1
        self.history.append(self.x_to_bbox_func(self.kf.x))
        return self.history[-1]

    def get_state(self):
        """
        Returns the current bounding box estimate.
        """
        return self.x_to_bbox_func(self.kf.x)

    def get_behavior(self):
        return self.behavior

    def mahalanobis(self, bbox):
        """Should be run after a predict() call for accuracy."""
        return self.kf.md_for_measurement(self.bbox_to_z_func(bbox))


"""
    We support multiple ways for association cost calculation, by default
    we use IoU. GIoU may have better performance in some situations. We note 
    that we hardly normalize the cost by all methods to (0,1) which may not be 
    the best practice.
"""
ASSO_FUNCS = {
    "iou": iou_batch,
    "giou": giou_batch,
    "ciou": ciou_batch,
    "diou": diou_batch,
    "ct_dist": ct_dist,
}


class DeepOCSort(object):
    def __init__(
        self,
        models,
        det_thresh=0.3,
        reid_thresh=0.5,
        max_age=30,
        min_hits=3,
        iou_threshold=0.3,
        delta_t=3,
        asso_func="iou",
        inertia=0.2,
        w_association_emb=0.75,
        alpha_fixed_emb=0.95,
        aw_param=0.5,
        embedding_off=False,
        cmc_off=False,
        aw_off=False,
        **kwargs,
    ):
        """
        Sets key parameters for SORT
        """
        self.det_thresh = det_thresh
        self.reid_thresh = reid_thresh
        self.max_age = max_age
        self.min_hits = min_hits
        self.iou_threshold = iou_threshold
        self.delta_t = delta_t
        self.asso_func = ASSO_FUNCS[asso_func]
        self.inertia = inertia
        self.w_association_emb = w_association_emb
        self.alpha_fixed_emb = alpha_fixed_emb
        self.aw_param = aw_param
        self.embedding_off = embedding_off
        self.cmc_off = cmc_off
        self.aw_off = aw_off

        self.trackers = []
        self.frame_count = 0
        KalmanBoxTracker.count = 0

        self.models = models
        self.cmc = CameraMotionCompensation()

    def update(self, dets, img, masks, tag="blub"):
        """
        Params:
          dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
        Requires: this method must be called once for each frame even with empty detections (use np.empty((0, 5)) for frames without detections).
        Returns the a similar array, where the last column is the object ID.
        NOTE: The number of objects returned may differ from the number of detections provided.
        """

        assert isinstance(
            dets, np.ndarray
        ), f"Unsupported 'dets' input format '{type(dets)}', valid format is np.ndarray"
        assert isinstance(
            img, np.ndarray
        ), f"Unsupported 'img' input format '{type(img)}', valid format is np.ndarray"
        assert (
            len(dets.shape) == 2
        ), f"Unsupported 'dets' dimensions, valid number of dimensions is two"
        assert (
            dets.shape[1] == 6
        ), f"Unsupported 'dets' 2nd dimension lenght, valid lenghts is 6"

        raw_dets = dets[:, 0:6]
        raw_scores = dets[:, 4]
        
        remain_inds = raw_scores > self.det_thresh
        dets = raw_dets[remain_inds]
        masks = masks[remain_inds]
        scores = raw_scores[remain_inds]

        # Embedding
        if self.embedding_off or dets.shape[0] == 0:
            dets_embs = np.ones((dets.shape[0], 1))
        else:
            dets_embs, behaviors = self._get_features(dets[:, :4], img, masks)

        # CMC
        if not self.cmc_off:
            transform = self.cmc.compute_affine(img, dets[:, :4], tag)
            for trk in self.trackers:
                trk.apply_affine_correction(transform)

        trust = (dets[:, 4] - self.det_thresh) / (1 - self.det_thresh)
        af = self.alpha_fixed_emb
        # From [self.alpha_fixed_emb, 1], goes to 1 as detector is less confident
        dets_alpha = af + (1 - af) * (1 - trust)

        # get predicted locations from existing trackers.
        trks = np.zeros((len(self.trackers), 5))
        to_del, trk_embs, trk_embs_raw = [], [], []
        for t, trk in enumerate(trks):
            pos = self.trackers[t].predict()[0]
            trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
            if np.any(np.isnan(pos)):
                to_del.append(t)
            else:
                trk_embs.append(self.trackers[t].get_emb())
                trk_embs_raw.append(self.trackers[t].get_emb_raw())
        trks = np.ma.compress_rows(np.ma.masked_invalid(trks))

        if len(trk_embs) > 0:
            trk_embs = np.vstack(trk_embs)
            trk_embs_raw = np.vstack(trk_embs_raw)
        else:
            trk_embs = np.array(trk_embs)
            trk_embs_raw = np.array(trk_embs_raw)

        for t in reversed(to_del):
            self.trackers.pop(t)

        velocities = np.array(
            [
                trk.velocity if trk.velocity is not None else np.array((0, 0))
                for trk in self.trackers
            ]
        )
        last_boxes = np.array([trk.last_observation for trk in self.trackers])
        k_observations = np.array(
            [
                k_previous_obs(trk.observations, trk.age, self.delta_t)
                for trk in self.trackers
            ]
        )

        """
            First round of association
        """
        # (M detections X N tracks, final score)
        if self.embedding_off or dets.shape[0] == 0 or trk_embs.shape[0] == 0:
            stage1_emb_cost = None
        else:
            stage1_emb_cost = dets_embs @ trk_embs.T
        matched, unmatched_dets, unmatched_trks = associate(
            dets,
            trks,
            self.iou_threshold,
            velocities,
            k_observations,
            self.inertia,
            stage1_emb_cost,
            self.w_association_emb,
            self.aw_off,
            self.aw_param,
        )
        for m in matched:
            self.trackers[m[1]].update(dets[m[0], :5], dets[m[0], 5])
            self.trackers[m[1]].update_emb(dets_embs[m[0]], alpha=dets_alpha[m[0]])
            self.trackers[m[1]].update_emb_raw(dets_embs[m[0]], scores[m[0]])
            self.trackers[m[1]].update_behavior(behaviors[m[0]])

        """
            Second round of associaton by OCR
        """
        if unmatched_dets.shape[0] > 0 and unmatched_trks.shape[0] > 0:
            left_dets = dets[unmatched_dets]
            left_dets_embs = dets_embs[unmatched_dets]
            left_trks = last_boxes[unmatched_trks]
            left_trks_embs = trk_embs[unmatched_trks]

            iou_left = self.asso_func(left_dets, left_trks)
            # TODO: is better without this
            emb_cost_left = left_dets_embs @ left_trks_embs.T
            if self.embedding_off:
                emb_cost_left = np.zeros_like(emb_cost_left)
            iou_left = np.array(iou_left)
            if iou_left.max() > self.iou_threshold:
                """
                NOTE: by using a lower threshold, e.g., self.iou_threshold - 0.1, you may
                get a higher performance especially on MOT17/MOT20 datasets. But we keep it
                uniform here for simplicity
                """
                rematched_indices = linear_assignment(-iou_left)
                to_remove_det_indices, to_remove_trk_indices = [], []
                for m in rematched_indices:
                    det_ind, trk_ind = unmatched_dets[m[0]], unmatched_trks[m[1]]
                    if iou_left[m[0], m[1]] < self.iou_threshold:
                        continue
                    self.trackers[trk_ind].update(dets[det_ind, :5], dets[det_ind, 5])
                    self.trackers[trk_ind].update_emb(
                        dets_embs[det_ind], alpha=dets_alpha[det_ind]
                    )
                    self.trackers[trk_ind].update_emb_raw(
                        dets_embs[det_ind], scores[det_ind]
                    )
                    self.trackers[trk_ind].update_behavior(behaviors[det_ind])
                    to_remove_det_indices.append(det_ind)
                    to_remove_trk_indices.append(trk_ind)
                unmatched_dets = np.setdiff1d(
                    unmatched_dets, np.array(to_remove_det_indices)
                )
                unmatched_trks = np.setdiff1d(
                    unmatched_trks, np.array(to_remove_trk_indices)
                )

        """
            Third round of association by comparing raw feature embedding
        """
        if unmatched_dets.shape[0] > 0 and unmatched_trks.shape[0] > 0:
            left_dets_embs = dets_embs[unmatched_dets]
            left_trks_embs = trk_embs_raw[unmatched_trks]

            distances = torchreid.metrics.compute_distance_matrix(
                left_dets_embs, torch.Tensor(left_trks_embs), metric="cosine"
            )

            if distances.min() < self.reid_thresh:
                rematched_indices = linear_assignment(distances.numpy())
                to_remove_det_indices, to_remove_trk_indices = [], []
                for m in rematched_indices:
                    det_ind, trk_ind = unmatched_dets[m[0]], unmatched_trks[m[1]]
                    if distances[m[0], m[1]] > self.reid_thresh:
                        continue
                    self.trackers[trk_ind].update(dets[det_ind, :5], dets[det_ind, 5])
                    self.trackers[trk_ind].update_emb(
                        dets_embs[det_ind], alpha=dets_alpha[det_ind]
                    )
                    self.trackers[trk_ind].update_emb_raw(
                        dets_embs[det_ind], scores[det_ind]
                    )
                    self.trackers[trk_ind].update_behavior(behaviors[det_ind])
                    to_remove_det_indices.append(det_ind)
                    to_remove_trk_indices.append(trk_ind)
                unmatched_dets = np.setdiff1d(
                    unmatched_dets, np.array(to_remove_det_indices)
                )
                unmatched_trks = np.setdiff1d(
                    unmatched_trks, np.array(to_remove_trk_indices)
                )

        for m in unmatched_trks:
            self.trackers[m].update(None, None)

        # create and initialise new trackers for unmatched detections
        for i in unmatched_dets:
            trk = KalmanBoxTracker(
                dets[i, :5],
                dets[i, 5],
                behaviors[i],
                delta_t=self.delta_t,
                emb=dets_embs[i],
                alpha=dets_alpha[i],
            )
            self.trackers.append(trk)
        
        i = len(self.trackers)
        ret, emb_ret, behavior_ret = [], [], []
        for trk in reversed(self.trackers):
            if trk.last_observation.sum() < 0:
                d = trk.get_state()[0]
            else:
                """
                this is optional to use the recent observation or the kalman filter prediction,
                we didn't notice significant difference here
                """
                d = trk.last_observation[:4]
            if (trk.time_since_update < 1) and (
                trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits
            ):
                # +1 as MOT benchmark requires positive
                emb_ret.append(trk.get_emb_raw().numpy())
                behavior_ret.append(trk.get_behavior())
                ret.append(
                    np.concatenate((d, [trk.id + 1], [trk.conf], [trk.cls])).reshape(
                        1, -1
                    )
                )
            i -= 1
            # remove dead tracklet
            if trk.time_since_update > self.max_age:
                self.trackers.pop(i)
        if len(ret) > 0:
            return np.concatenate(ret), np.array(behavior_ret), np.array(emb_ret)
        return np.empty((0, 5)), np.array([]), np.array([])

    @torch.no_grad()
    def _get_features(self, bbox_xyxy, ori_img, masks):
        im_crops = []
        behaviors = []
        for i, box in enumerate(bbox_xyxy):
            x1, y1, x2, y2 = box.astype(int)
            im_seg = np.zeros_like(ori_img)
            im_seg[masks[i]] = ori_img[masks[i]]
            im_seg = im_seg[y1:y2, x1:x2]
            im_crops.append(im_seg)

            emotion, gender, posture = self.models.classify_emotion_gender_posture(
                im_seg
            )
            behaviors.append([emotion, gender, posture])
        if im_crops:
            return self.models.reid_model(im_crops).cpu(), np.array(behaviors)
        return np.array([]), np.array([])