Source code for gluoncv.model_zoo.smot.utils
"""
Utility functions for SMOT: Single-Shot Multi Object Tracking
https://arxiv.org/abs/2010.16031
"""
import time
import logging
from contextlib import contextmanager
import mxnet as mx
import numpy as np
[docs]def timeit(method):
"""
The timing decorator to wrap the functions
"""
def timed(*args, **kw):
ts = time.time()
result = method(*args, **kw)
te = time.time()
logging.info("{} runtime: {:.04f} msec".format(method.__name__, (te - ts) * 1000))
return result
return timed
@contextmanager
def timeit_context(name):
startTime = time.time()
yield
elapsedTime = time.time() - startTime
logging.info('[{}] runtime {:.03f} msec'.format(name, elapsedTime * 1000))
def mxnet_frame_preprocessing(image, base_size, ratio, mean, std, ctx):
"""
Parameters
----------
image
base_size
ratio: aspect ratio
mean
std
ctx
Returns
-------
"""
image_tensor = mx.nd.array(image, ctx=ctx, dtype=np.uint8)
float_image = (image_tensor.astype(np.float32) / 255. - mean.reshape((1, 1, 3))) / std.reshape((1, 1, 3))
trans_image = float_image.transpose((2, 0, 1))
out_w, out_h = int(base_size * ratio), int(base_size)
in_h, in_w, _ = image.shape
pad_w, pad_h = int(max(in_w, in_h * ratio)), int(max(in_h, in_w / ratio))
pb_w, pb_h = (pad_w - in_w) // 2, (pad_h - in_h) // 2
# do padding
padded_image = mx.nd.zeros((1, 3, pad_h, pad_w), ctx=ctx)
padded_image[0, :, pb_h: pb_h + in_h, pb_w: pb_w + in_w] = trans_image
# do resizing
resize_image = mx.nd.contrib.BilinearResize2D(padded_image,
height=out_h, width=out_w)
return resize_image, pad_w, pad_h, (pb_w, pb_h, pad_w, pad_h)
def remap_bboxes(bboxes, padded_w, padded_h, expand, data_shape, ratio):
"""
Remap bboxes in (x0, y0, x1, y1) format into the input image space
Parameters
----------
bboxes
padded_w
padded_h
expand
Returns
-------
"""
bboxes[:, 0] *= padded_w / (data_shape * ratio)
bboxes[:, 1] *= padded_h / data_shape
bboxes[:, 2] *= padded_w / (data_shape * ratio)
bboxes[:, 3] *= padded_h / data_shape
bboxes[:, 0] -= expand[0]
bboxes[:, 1] -= expand[1]
bboxes[:, 2] -= expand[0]
bboxes[:, 3] -= expand[1]
return bboxes
class TrackState:
"""
States of the track.
The track follows the simple state machine as below:
Active: time_since_update always set to 1
1. If confidence < keep_alive_threshold, goto Missing
2. If the track is suppressed in track NMS, goto Missing
Missing: every timestep the missing track increment time_since_update by one
1. If the track is updated again, goto Active
2. If time_since_update > max_missing, goto Deleted
Deleted: This is an absorbing state
"""
Active = 1
Missing = 2
Deleted = 3
[docs]class Track:
"""
This class represents a track/tracklet used in the SMOT Tracker
It has the following properties
*******************************************************
mean: 4-tuple representing the (x0, y0, x1, y1) as the current state (location) of the tracked object
track_id: the numerical id of the track
age: the number of timesteps since its first occurrence
time_since_update: number of time-steps since the last update of the its location
state: the state of the track, can be one in `TrackState`
confidence_score: tracking_confidence at the current timestep
source: a tuple of (anchor_indices, anchor_weights)
attributes: np.ndarray of additional attributes of the object
*******************************************************
It also has these configs
keep_alive_thresh: the minimal tracking/detection confidence to keep the track in Active state
max_missing: the maximal timesteps we will keep searching for this track when missing before we mark it as deleted
*******************************************************
"""
def __init__(self, mean, track_id, source, keep_alive_thresh=0.1, max_missing=30,
attributes=None, class_id=0, linked_id=None):
self.mean = mean
self.track_id = track_id
self.hits = 1
self.age = 1
self.time_since_update = 0
self.state = TrackState.Active
self.confidence_score = 1.
self.matched = False
self.match_score = 999
self.current_detection_index = None
self.max_missing = max_missing
self.next_frame_anchor_id = -1
self.keep_alive_thresh = keep_alive_thresh
self.attributes = attributes
self.source = source
self.class_id = class_id
self.linked_id = linked_id
@property
def display_id(self):
if self.linked_id is not None:
return self.linked_id
else:
return self.track_id
@property
def linkable(self):
return self.linked_id is None
def link_to(self, track):
self.linked_id = track.display_id
[docs] def predict(self, motion_model=None):
"""
Parameters
----------
motion_model : if not None, predict the motion of this track given its history
"""
if motion_model:
self.mean, self.covariance = motion_model(self.mean, self.covariance)
self.age += 1
self.time_since_update += 1
# If the track has long been missing, mark it as deleted
if self.time_since_update > self.max_missing:
self.state = TrackState.Deleted
[docs] def update(self, bbx, source=None, attributes=None):
"""
Update the state of the track. We override the predicted track position.
Updating the track will keep or flip its state as Active
If the confidence of detection is below the keep_alive_threshold, we will mark this track as missed.
----------
bbx : new detection location of this object
attributes: some useful attributes of this object at this frame, e.g. landmarks
"""
self.mean = bbx[:4]
self.confidence_score = bbx[4]
self.time_since_update = 0
self.state = TrackState.Active
self.attributes = attributes
self.source = source
if self.confidence_score < self.keep_alive_thresh:
self.mark_missed()
[docs] def mark_missed(self):
"""Mark this track as missed (no association at the current time step).
"""
# only an active track can be marked as missed
if self.state == TrackState.Active:
self.state = TrackState.Missing
[docs] def is_mising(self):
"""Returns True if this track is tentative (unconfirmed).
"""
return self.state == TrackState.Missing
[docs] def is_active(self):
"""Returns True if this track is confirmed."""
return self.state == TrackState.Active
[docs] def is_deleted(self):
"""Returns True if this track is dead and should be deleted."""
return self.state == TrackState.Deleted