|
|
|
|
|
import copy
|
|
import numpy as np
|
|
from typing import Dict
|
|
import torch
|
|
from scipy.optimize import linear_sum_assignment
|
|
|
|
from detectron2.config import configurable
|
|
from detectron2.structures import Boxes, Instances
|
|
|
|
from ..config.config import CfgNode as CfgNode_
|
|
from .base_tracker import BaseTracker
|
|
|
|
|
|
class BaseHungarianTracker(BaseTracker):
|
|
"""
|
|
A base class for all Hungarian trackers
|
|
"""
|
|
|
|
@configurable
|
|
def __init__(
|
|
self,
|
|
video_height: int,
|
|
video_width: int,
|
|
max_num_instances: int = 200,
|
|
max_lost_frame_count: int = 0,
|
|
min_box_rel_dim: float = 0.02,
|
|
min_instance_period: int = 1,
|
|
**kwargs
|
|
):
|
|
"""
|
|
Args:
|
|
video_height: height the video frame
|
|
video_width: width of the video frame
|
|
max_num_instances: maximum number of id allowed to be tracked
|
|
max_lost_frame_count: maximum number of frame an id can lost tracking
|
|
exceed this number, an id is considered as lost
|
|
forever
|
|
min_box_rel_dim: a percentage, smaller than this dimension, a bbox is
|
|
removed from tracking
|
|
min_instance_period: an instance will be shown after this number of period
|
|
since its first showing up in the video
|
|
"""
|
|
super().__init__(**kwargs)
|
|
self._video_height = video_height
|
|
self._video_width = video_width
|
|
self._max_num_instances = max_num_instances
|
|
self._max_lost_frame_count = max_lost_frame_count
|
|
self._min_box_rel_dim = min_box_rel_dim
|
|
self._min_instance_period = min_instance_period
|
|
|
|
@classmethod
|
|
def from_config(cls, cfg: CfgNode_) -> Dict:
|
|
raise NotImplementedError("Calling HungarianTracker::from_config")
|
|
|
|
def build_cost_matrix(self, instances: Instances, prev_instances: Instances) -> np.ndarray:
|
|
raise NotImplementedError("Calling HungarianTracker::build_matrix")
|
|
|
|
def update(self, instances: Instances) -> Instances:
|
|
if instances.has("pred_keypoints"):
|
|
raise NotImplementedError("Need to add support for keypoints")
|
|
instances = self._initialize_extra_fields(instances)
|
|
if self._prev_instances is not None:
|
|
self._untracked_prev_idx = set(range(len(self._prev_instances)))
|
|
cost_matrix = self.build_cost_matrix(instances, self._prev_instances)
|
|
matched_idx, matched_prev_idx = linear_sum_assignment(cost_matrix)
|
|
instances = self._process_matched_idx(instances, matched_idx, matched_prev_idx)
|
|
instances = self._process_unmatched_idx(instances, matched_idx)
|
|
instances = self._process_unmatched_prev_idx(instances, matched_prev_idx)
|
|
self._prev_instances = copy.deepcopy(instances)
|
|
return instances
|
|
|
|
def _initialize_extra_fields(self, instances: Instances) -> Instances:
|
|
"""
|
|
If input instances don't have ID, ID_period, lost_frame_count fields,
|
|
this method is used to initialize these fields.
|
|
|
|
Args:
|
|
instances: D2 Instances, for predictions of the current frame
|
|
Return:
|
|
D2 Instances with extra fields added
|
|
"""
|
|
if not instances.has("ID"):
|
|
instances.set("ID", [None] * len(instances))
|
|
if not instances.has("ID_period"):
|
|
instances.set("ID_period", [None] * len(instances))
|
|
if not instances.has("lost_frame_count"):
|
|
instances.set("lost_frame_count", [None] * len(instances))
|
|
if self._prev_instances is None:
|
|
instances.ID = list(range(len(instances)))
|
|
self._id_count += len(instances)
|
|
instances.ID_period = [1] * len(instances)
|
|
instances.lost_frame_count = [0] * len(instances)
|
|
return instances
|
|
|
|
def _process_matched_idx(
|
|
self, instances: Instances, matched_idx: np.ndarray, matched_prev_idx: np.ndarray
|
|
) -> Instances:
|
|
assert matched_idx.size == matched_prev_idx.size
|
|
for i in range(matched_idx.size):
|
|
instances.ID[matched_idx[i]] = self._prev_instances.ID[matched_prev_idx[i]]
|
|
instances.ID_period[matched_idx[i]] = (
|
|
self._prev_instances.ID_period[matched_prev_idx[i]] + 1
|
|
)
|
|
instances.lost_frame_count[matched_idx[i]] = 0
|
|
return instances
|
|
|
|
def _process_unmatched_idx(self, instances: Instances, matched_idx: np.ndarray) -> Instances:
|
|
untracked_idx = set(range(len(instances))).difference(set(matched_idx))
|
|
for idx in untracked_idx:
|
|
instances.ID[idx] = self._id_count
|
|
self._id_count += 1
|
|
instances.ID_period[idx] = 1
|
|
instances.lost_frame_count[idx] = 0
|
|
return instances
|
|
|
|
def _process_unmatched_prev_idx(
|
|
self, instances: Instances, matched_prev_idx: np.ndarray
|
|
) -> Instances:
|
|
untracked_instances = Instances(
|
|
image_size=instances.image_size,
|
|
pred_boxes=[],
|
|
pred_masks=[],
|
|
pred_classes=[],
|
|
scores=[],
|
|
ID=[],
|
|
ID_period=[],
|
|
lost_frame_count=[],
|
|
)
|
|
prev_bboxes = list(self._prev_instances.pred_boxes)
|
|
prev_classes = list(self._prev_instances.pred_classes)
|
|
prev_scores = list(self._prev_instances.scores)
|
|
prev_ID_period = self._prev_instances.ID_period
|
|
if instances.has("pred_masks"):
|
|
prev_masks = list(self._prev_instances.pred_masks)
|
|
untracked_prev_idx = set(range(len(self._prev_instances))).difference(set(matched_prev_idx))
|
|
for idx in untracked_prev_idx:
|
|
x_left, y_top, x_right, y_bot = prev_bboxes[idx]
|
|
if (
|
|
(1.0 * (x_right - x_left) / self._video_width < self._min_box_rel_dim)
|
|
or (1.0 * (y_bot - y_top) / self._video_height < self._min_box_rel_dim)
|
|
or self._prev_instances.lost_frame_count[idx] >= self._max_lost_frame_count
|
|
or prev_ID_period[idx] <= self._min_instance_period
|
|
):
|
|
continue
|
|
untracked_instances.pred_boxes.append(list(prev_bboxes[idx].numpy()))
|
|
untracked_instances.pred_classes.append(int(prev_classes[idx]))
|
|
untracked_instances.scores.append(float(prev_scores[idx]))
|
|
untracked_instances.ID.append(self._prev_instances.ID[idx])
|
|
untracked_instances.ID_period.append(self._prev_instances.ID_period[idx])
|
|
untracked_instances.lost_frame_count.append(
|
|
self._prev_instances.lost_frame_count[idx] + 1
|
|
)
|
|
if instances.has("pred_masks"):
|
|
untracked_instances.pred_masks.append(prev_masks[idx].numpy().astype(np.uint8))
|
|
|
|
untracked_instances.pred_boxes = Boxes(torch.FloatTensor(untracked_instances.pred_boxes))
|
|
untracked_instances.pred_classes = torch.IntTensor(untracked_instances.pred_classes)
|
|
untracked_instances.scores = torch.FloatTensor(untracked_instances.scores)
|
|
if instances.has("pred_masks"):
|
|
untracked_instances.pred_masks = torch.IntTensor(untracked_instances.pred_masks)
|
|
else:
|
|
untracked_instances.remove("pred_masks")
|
|
|
|
return Instances.cat(
|
|
[
|
|
instances,
|
|
untracked_instances,
|
|
]
|
|
)
|
|
|
