| | |
| | from abc import ABCMeta, abstractmethod |
| | from typing import List, Optional, Tuple |
| |
|
| | import torch |
| | import torch.nn.functional as F |
| | from addict import Dict |
| |
|
| |
|
| | class BaseTracker(metaclass=ABCMeta): |
| | """Base tracker model. |
| | |
| | Args: |
| | momentums (dict[str:float], optional): Momentums to update the buffers. |
| | The `str` indicates the name of the buffer while the `float` |
| | indicates the momentum. Defaults to None. |
| | num_frames_retain (int, optional). If a track is disappeared more than |
| | `num_frames_retain` frames, it will be deleted in the memo. |
| | Defaults to 10. |
| | """ |
| |
|
| | def __init__(self, |
| | momentums: Optional[dict] = None, |
| | num_frames_retain: int = 10) -> None: |
| | super().__init__() |
| | if momentums is not None: |
| | assert isinstance(momentums, dict), 'momentums must be a dict' |
| | self.momentums = momentums |
| | self.num_frames_retain = num_frames_retain |
| |
|
| | self.reset() |
| |
|
| | def reset(self) -> None: |
| | """Reset the buffer of the tracker.""" |
| | self.num_tracks = 0 |
| | self.tracks = dict() |
| |
|
| | @property |
| | def empty(self) -> bool: |
| | """Whether the buffer is empty or not.""" |
| | return False if self.tracks else True |
| |
|
| | @property |
| | def ids(self) -> List[dict]: |
| | """All ids in the tracker.""" |
| | return list(self.tracks.keys()) |
| |
|
| | @property |
| | def with_reid(self) -> bool: |
| | """bool: whether the framework has a reid model""" |
| | return hasattr(self, 'reid') and self.reid is not None |
| |
|
| | def update(self, **kwargs) -> None: |
| | """Update the tracker. |
| | |
| | Args: |
| | kwargs (dict[str: Tensor | int]): The `str` indicates the |
| | name of the input variable. `ids` and `frame_ids` are |
| | obligatory in the keys. |
| | """ |
| | memo_items = [k for k, v in kwargs.items() if v is not None] |
| | rm_items = [k for k in kwargs.keys() if k not in memo_items] |
| | for item in rm_items: |
| | kwargs.pop(item) |
| | if not hasattr(self, 'memo_items'): |
| | self.memo_items = memo_items |
| | else: |
| | assert memo_items == self.memo_items |
| |
|
| | assert 'ids' in memo_items |
| | num_objs = len(kwargs['ids']) |
| | id_indice = memo_items.index('ids') |
| | assert 'frame_ids' in memo_items |
| | frame_id = int(kwargs['frame_ids']) |
| | if isinstance(kwargs['frame_ids'], int): |
| | kwargs['frame_ids'] = torch.tensor([kwargs['frame_ids']] * |
| | num_objs) |
| | |
| | for k, v in kwargs.items(): |
| | if len(v) != num_objs: |
| | raise ValueError('kwargs value must both equal') |
| |
|
| | for obj in zip(*kwargs.values()): |
| | id = int(obj[id_indice]) |
| | if id in self.tracks: |
| | self.update_track(id, obj) |
| | else: |
| | self.init_track(id, obj) |
| |
|
| | self.pop_invalid_tracks(frame_id) |
| |
|
| | def pop_invalid_tracks(self, frame_id: int) -> None: |
| | """Pop out invalid tracks.""" |
| | invalid_ids = [] |
| | for k, v in self.tracks.items(): |
| | if frame_id - v['frame_ids'][-1] >= self.num_frames_retain: |
| | invalid_ids.append(k) |
| | for invalid_id in invalid_ids: |
| | self.tracks.pop(invalid_id) |
| |
|
| | def update_track(self, id: int, obj: Tuple[torch.Tensor]): |
| | """Update a track.""" |
| | for k, v in zip(self.memo_items, obj): |
| | v = v[None] |
| | if self.momentums is not None and k in self.momentums: |
| | m = self.momentums[k] |
| | self.tracks[id][k] = (1 - m) * self.tracks[id][k] + m * v |
| | else: |
| | self.tracks[id][k].append(v) |
| |
|
| | def init_track(self, id: int, obj: Tuple[torch.Tensor]): |
| | """Initialize a track.""" |
| | self.tracks[id] = Dict() |
| | for k, v in zip(self.memo_items, obj): |
| | v = v[None] |
| | if self.momentums is not None and k in self.momentums: |
| | self.tracks[id][k] = v |
| | else: |
| | self.tracks[id][k] = [v] |
| |
|
| | @property |
| | def memo(self) -> dict: |
| | """Return all buffers in the tracker.""" |
| | outs = Dict() |
| | for k in self.memo_items: |
| | outs[k] = [] |
| |
|
| | for id, objs in self.tracks.items(): |
| | for k, v in objs.items(): |
| | if k not in outs: |
| | continue |
| | if self.momentums is not None and k in self.momentums: |
| | v = v |
| | else: |
| | v = v[-1] |
| | outs[k].append(v) |
| |
|
| | for k, v in outs.items(): |
| | outs[k] = torch.cat(v, dim=0) |
| | return outs |
| |
|
| | def get(self, |
| | item: str, |
| | ids: Optional[list] = None, |
| | num_samples: Optional[int] = None, |
| | behavior: Optional[str] = None) -> torch.Tensor: |
| | """Get the buffer of a specific item. |
| | |
| | Args: |
| | item (str): The demanded item. |
| | ids (list[int], optional): The demanded ids. Defaults to None. |
| | num_samples (int, optional): Number of samples to calculate the |
| | results. Defaults to None. |
| | behavior (str, optional): Behavior to calculate the results. |
| | Options are `mean` | None. Defaults to None. |
| | |
| | Returns: |
| | Tensor: The results of the demanded item. |
| | """ |
| | if ids is None: |
| | ids = self.ids |
| |
|
| | outs = [] |
| | for id in ids: |
| | out = self.tracks[id][item] |
| | if isinstance(out, list): |
| | if num_samples is not None: |
| | out = out[-num_samples:] |
| | out = torch.cat(out, dim=0) |
| | if behavior == 'mean': |
| | out = out.mean(dim=0, keepdim=True) |
| | elif behavior is None: |
| | out = out[None] |
| | else: |
| | raise NotImplementedError() |
| | else: |
| | out = out[-1] |
| | outs.append(out) |
| | return torch.cat(outs, dim=0) |
| |
|
| | @abstractmethod |
| | def track(self, *args, **kwargs): |
| | """Tracking forward function.""" |
| | pass |
| |
|
| | def crop_imgs(self, |
| | img: torch.Tensor, |
| | meta_info: dict, |
| | bboxes: torch.Tensor, |
| | rescale: bool = False) -> torch.Tensor: |
| | """Crop the images according to some bounding boxes. Typically for re- |
| | identification sub-module. |
| | |
| | Args: |
| | img (Tensor): of shape (T, C, H, W) encoding input image. |
| | Typically these should be mean centered and std scaled. |
| | meta_info (dict): image information dict where each dict |
| | has: 'img_shape', 'scale_factor', 'flip', and may also contain |
| | 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. |
| | bboxes (Tensor): of shape (N, 4) or (N, 5). |
| | rescale (bool, optional): If True, the bounding boxes should be |
| | rescaled to fit the scale of the image. Defaults to False. |
| | |
| | Returns: |
| | Tensor: Image tensor of shape (T, C, H, W). |
| | """ |
| | h, w = meta_info['img_shape'] |
| | img = img[:, :, :h, :w] |
| | if rescale: |
| | factor_x, factor_y = meta_info['scale_factor'] |
| | bboxes[:, :4] *= torch.tensor( |
| | [factor_x, factor_y, factor_x, factor_y]).to(bboxes.device) |
| | bboxes[:, 0] = torch.clamp(bboxes[:, 0], min=0, max=w - 1) |
| | bboxes[:, 1] = torch.clamp(bboxes[:, 1], min=0, max=h - 1) |
| | bboxes[:, 2] = torch.clamp(bboxes[:, 2], min=1, max=w) |
| | bboxes[:, 3] = torch.clamp(bboxes[:, 3], min=1, max=h) |
| |
|
| | crop_imgs = [] |
| | for bbox in bboxes: |
| | x1, y1, x2, y2 = map(int, bbox) |
| | if x2 <= x1: |
| | x2 = x1 + 1 |
| | if y2 <= y1: |
| | y2 = y1 + 1 |
| | crop_img = img[:, :, y1:y2, x1:x2] |
| | if self.reid.get('img_scale', False): |
| | crop_img = F.interpolate( |
| | crop_img, |
| | size=self.reid['img_scale'], |
| | mode='bilinear', |
| | align_corners=False) |
| | crop_imgs.append(crop_img) |
| |
|
| | if len(crop_imgs) > 0: |
| | return torch.cat(crop_imgs, dim=0) |
| | elif self.reid.get('img_scale', False): |
| | _h, _w = self.reid['img_scale'] |
| | return img.new_zeros((0, 3, _h, _w)) |
| | else: |
| | return img.new_zeros((0, 3, h, w)) |
| |
|
