import torch
from addict import Dict
from ..builder import MODELS, build_detector
from .base import BaseVideoDetector
[docs]@MODELS.register_module()
class SELSA(BaseVideoDetector):
"""Sequence Level Semantics Aggregation for Video Object Detection.
This video object detector is the implementation of `SELSA
<https://arxiv.org/abs/1907.06390>`_.
"""
def __init__(self,
detector,
pretrains=None,
frozen_modules=None,
train_cfg=None,
test_cfg=None):
super(SELSA, self).__init__()
self.detector = build_detector(detector)
assert hasattr(self.detector, 'roi_head'), \
'selsa video detector only supports two stage detector'
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrains)
if frozen_modules is not None:
self.freeze_module(frozen_modules)
[docs] def init_weights(self, pretrain):
"""Initialize the weights of modules in video object detector.
Args:
pretrained (dict): Path to pre-trained weights.
"""
if pretrain is None:
pretrain = dict()
assert isinstance(pretrain, dict), '`pretrain` must be a dict.'
if self.with_detector and pretrain.get('detector', False):
self.init_module('detector', pretrain['detector'])
if self.with_motion:
self.init_module('motion', pretrain.get('motion', None))
[docs] def forward_train(self,
img,
img_metas,
gt_bboxes,
gt_labels,
ref_img,
ref_img_metas,
ref_gt_bboxes,
ref_gt_labels,
gt_instance_ids=None,
gt_bboxes_ignore=None,
gt_masks=None,
proposals=None,
ref_gt_instance_ids=None,
ref_gt_bboxes_ignore=None,
ref_gt_masks=None,
ref_proposals=None,
**kwargs):
"""
Args:
img (Tensor): of shape (N, C, H, W) encoding input images.
Typically these should be mean centered and std scaled.
img_metas (list[dict]): list of image info dict where each dict
has: 'img_shape', 'scale_factor', 'flip', and may also contain
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
For details on the values of these keys see
`mmtrack/datasets/pipelines/formatting.py:VideoCollect`.
gt_bboxes (list[Tensor]): Ground truth bboxes for each image with
shape (num_gts, 4) in [tl_x, tl_y, br_x, br_y] format.
gt_labels (list[Tensor]): class indices corresponding to each box.
ref_img (Tensor): of shape (N, 2, C, H, W) encoding input images.
Typically these should be mean centered and std scaled.
2 denotes there is two reference images for each input image.
ref_img_metas (list[list[dict]]): The first list only has one
element. The second list contains reference 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'. For details on the values of these keys see
`mmtrack/datasets/pipelines/formatting.py:VideoCollect`.
ref_gt_bboxes (list[Tensor]): The list only has one Tensor. The
Tensor contains ground truth bboxes for each reference image
with shape (num_all_ref_gts, 5) in
[ref_img_id, tl_x, tl_y, br_x, br_y] format. The ref_img_id
start from 0, and denotes the id of reference image for each
key image.
ref_gt_labels (list[Tensor]): The list only has one Tensor. The
Tensor contains class indices corresponding to each reference
box with shape (num_all_ref_gts, 2) in
[ref_img_id, class_indice].
gt_instance_ids (None | list[Tensor]): specify the instance id for
each ground truth bbox.
gt_bboxes_ignore (None | list[Tensor]): specify which bounding
boxes can be ignored when computing the loss.
gt_masks (None | Tensor) : true segmentation masks for each box
used if the architecture supports a segmentation task.
proposals (None | Tensor) : override rpn proposals with custom
proposals. Use when `with_rpn` is False.
ref_gt_instance_ids (None | list[Tensor]): specify the instance id
for each ground truth bboxes of reference images.
ref_gt_bboxes_ignore (None | list[Tensor]): specify which bounding
boxes of reference images can be ignored when computing the
loss.
ref_gt_masks (None | Tensor) : True segmentation masks for each
box of reference image used if the architecture supports a
segmentation task.
ref_proposals (None | Tensor) : override rpn proposals with custom
proposals of reference images. Use when `with_rpn` is False.
Returns:
dict[str, Tensor]: a dictionary of loss components
"""
assert len(img) == 1, \
'selsa video detector only supports 1 batch size per gpu for now.'
all_imgs = torch.cat((img, ref_img[0]), dim=0)
all_x = self.detector.extract_feat(all_imgs)
x = []
ref_x = []
for i in range(len(all_x)):
x.append(all_x[i][[0]])
ref_x.append(all_x[i][1:])
losses = dict()
# RPN forward and loss
if self.detector.with_rpn:
proposal_cfg = self.detector.train_cfg.get(
'rpn_proposal', self.detector.test_cfg.rpn)
rpn_losses, proposal_list = self.detector.rpn_head.forward_train(
x,
img_metas,
gt_bboxes,
gt_labels=None,
gt_bboxes_ignore=gt_bboxes_ignore,
proposal_cfg=proposal_cfg)
losses.update(rpn_losses)
ref_proposals_list = self.detector.rpn_head.simple_test_rpn(
ref_x, ref_img_metas[0])
else:
proposal_list = proposals
ref_proposals_list = ref_proposals
roi_losses = self.detector.roi_head.forward_train(
x, ref_x, img_metas, proposal_list, ref_proposals_list, gt_bboxes,
gt_labels, gt_bboxes_ignore, gt_masks, **kwargs)
losses.update(roi_losses)
return losses
[docs] def simple_test(self,
img,
img_metas,
ref_img=None,
ref_img_metas=None,
proposals=None,
ref_proposals=None,
rescale=False):
"""Test without augmentation.
Args:
img (Tensor): of shape (1, C, H, W) encoding input image.
Typically these should be mean centered and std scaled.
img_metas (list[dict]): list of 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'. For details on the values of these keys see
`mmtrack/datasets/pipelines/formatting.py:VideoCollect`.
ref_img (list[Tensor] | None): The list only contains one Tensor
of shape (1, N, C, H, W) encoding input reference images.
Typically these should be mean centered and std scaled. N
denotes the number for reference images. There may be no
reference images in some cases.
ref_img_metas (list[list[list[dict]]] | None): The first and
second list only has one element. The third list contains
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'. For details on
the values of these keys see
`mmtrack/datasets/pipelines/formatting.py:VideoCollect`. There
may be no reference images in some cases.
proposals (None | Tensor): Override rpn proposals with custom
proposals. Use when `with_rpn` is False. Defaults to None.
rescale (bool): If False, then returned bboxes and masks will fit
the scale of img, otherwise, returned bboxes and masks
will fit the scale of original image shape. Defaults to False.
Returns:
dict[str : list(ndarray)]: The detection results.
"""
if ref_img is not None:
ref_img = ref_img[0]
if ref_img_metas is not None:
ref_img_metas = ref_img_metas[0]
x, img_metas, ref_x, ref_img_metas = self.extract_feats(
img, img_metas, ref_img, ref_img_metas)
if proposals is None:
proposal_list = self.detector.rpn_head.simple_test_rpn(
x, img_metas)
ref_proposals_list = self.detector.rpn_head.simple_test_rpn(
ref_x, ref_img_metas)
else:
proposal_list = proposals
ref_proposals_list = ref_proposals
outs = self.detector.roi_head.simple_test(
x,
ref_x,
proposal_list,
ref_proposals_list,
img_metas,
rescale=rescale)
results = dict()
results['bbox_results'] = outs[0]
if len(outs) == 2:
results['segm_results'] = outs[1]
return results
[docs] def aug_test(self, imgs, img_metas, **kwargs):
"""Test function with test time augmentation."""
raise NotImplementedError