Source code for gluoncv.model_zoo.smot.tracktors

"""
Tractor module for SMOT
"""
# pylint: disable=arguments-differ
from abc import ABC, abstractmethod
import numpy as np

import mxnet as mx

from .utils import timeit
from .general_detector import GeneralDetector


[docs]class BaseAnchorBasedTracktor(ABC): """ BaseAnchorBasedTracktor """
[docs] @abstractmethod def anchors(self): """ Returns the list of anchors used in this detector. ------- """ raise NotImplementedError
[docs] @abstractmethod def prepare_for_frame(self, frame): """ This method should run anything that needs to happen before the motion prediction. It can prepare the detector or even run the backbone feature extractions. It can also provide data to motion prediction Parameters ---------- frame: the frame data, the same as in the detect_and_track method Returns ------- motion_predict_data: optional data provided to motion prediction, if no data is provided, return None """ raise NotImplementedError
[docs] @abstractmethod def detect_and_track(self, frame, tracking_anchor_indices, tracking_anchor_weights, tracking_classes): """ Perform detection and tracking on the new frame Parameters ---------- frame: HxWx3 RGB image tracking_anchor_indices: NxM ndarray tracking_anchor_weights NxM ndarray tracking_classes: Nx1 ndarray of the class ids of the tracked object Returns detection_bounding_boxes: all detection results, in (x0, y0, x1, y1, confidence, cls) format detection_source: source anchor box indices for each detection tracking_boxes: all tracking results, in (x0, y0, x1, y1, confidence) format extract_info: extra information from the tracktor, e.g. landmarks, a dict ------- """ raise NotImplementedError
[docs] @abstractmethod def clean_up(self): """ Clean up after running one video Returns ------- """ raise NotImplementedError
[docs]class GluonSSDMultiClassTracktor(BaseAnchorBasedTracktor): """ Initiate a tracktor based on an object detetor. """ def __init__(self, gpu_id=0, detector_thresh=0.5, model_name="", use_pretrained=False, param_path="", data_shape=512): self.detector = GeneralDetector(gpu_id, aspect_ratio=1., data_shape=data_shape, model_name=model_name, use_pretrained=use_pretrained, param_path=param_path ) self._anchor_tensor = None self._detector_thresh = detector_thresh self._ctx = mx.gpu(gpu_id) self._dummy_ti = mx.nd.array([[0]], ctx=self._ctx)
[docs] def anchors(self): if self.detector.anchor_tensor is None: raise ValueError("anchor not initialized yet") return self.detector.anchor_tensor
[docs] def clean_up(self): pass
[docs] def prepare_for_frame(self, frame): return None
@timeit def detect_and_track(self, frame, tracking_anchor_indices, tracking_anchor_weights, tracking_object_classes): with_tracking = len(tracking_anchor_indices) > 0 if with_tracking: tracking_indices = mx.nd.array(tracking_anchor_indices, ctx=self._ctx) tracking_weights = mx.nd.array(tracking_anchor_weights, ctx=self._ctx) tracking_classes = mx.nd.array(tracking_object_classes.reshape((-1, 1)), ctx=self._ctx) else: tracking_classes = tracking_indices = tracking_weights = self._dummy_ti ids, scores, bboxes, voted_tracking_bboxes, detection_anchor_indices = \ self.detector.run_detection(frame, tracking_indices, tracking_weights, tracking_classes) valid_det_num = (scores > self._detector_thresh).sum().astype(int).asnumpy()[0] if valid_det_num > 0: valid_scores = scores[:valid_det_num] valid_bboxes = bboxes[:valid_det_num, :] valid_classes = ids[:valid_det_num, :] detection_output = mx.nd.concat(valid_bboxes, valid_scores, valid_classes, dim=1).asnumpy() anchor_indices_output = detection_anchor_indices[:valid_det_num, :] else: # no detection detection_output = np.array([]) anchor_indices_output = np.array([]) tracking_response = voted_tracking_bboxes.asnumpy() \ if with_tracking else np.array([]) return detection_output, anchor_indices_output, tracking_response, \ {}