def forward_dummy(self, x, proposals):
outs = ()
for i in range(self.num_stages):
start_bbox_type = self.bbox_head[i].start_bbox_type
if start_bbox_type == 'hbb':
rois = arb2roi([proposals], bbox_type='hbb')
else:
theta = proposals.new_zeros((proposals.size(0), 1))
proposal_theta = torch.cat([proposals, theta], dim=1)
rois = arb2roi([proposal_theta], bbox_type='obb')
bbox_results = self._bbox_forward(i, x, rois)
outs = outs + (bbox_results['cls_score'], bbox_results['bbox_pred'])
return outs
async def async_test_bboxes(self,
x,
img_metas,
proposals,
rcnn_test_cfg,
rescale=False,
bbox_semaphore=None,
global_lock=None):
"""Asynchronized test for box head without augmentation."""
rois = arb2roi(proposals, bbox_type=self.bbox_head.start_bbox_type)
roi_feats = self.bbox_roi_extractor(
x[:len(self.bbox_roi_extractor.featmap_strides)], rois)
if self.with_shared_head:
roi_feats = self.shared_head(roi_feats)
sleep_interval = rcnn_test_cfg.get('async_sleep_interval', 0.017)
async with completed(
__name__, 'bbox_head_forward',
sleep_interval=sleep_interval):
cls_score, bbox_pred = self.bbox_head(roi_feats)
img_shape = img_metas[0]['img_shape']
scale_factor = img_metas[0]['scale_factor']
det_bboxes, det_labels = self.bbox_head.get_bboxes(
rois,
cls_score,
bbox_pred,
img_shape,
scale_factor,
rescale=rescale,
cfg=rcnn_test_cfg)
return det_bboxes, det_labels
def simple_test_mask(self,
x,
img_metas,
det_bboxes,
det_labels,
rescale=False):
"""Simple test for mask head without augmentation."""
# image shape of the first image in the batch (only one)
ori_shape = img_metas[0]['ori_shape']
scale_factor = img_metas[0]['scale_factor']
if det_bboxes.shape[0] == 0:
segm_result = [[] for _ in range(self.mask_head.num_classes)]
else:
# if det_bboxes is rescaled to the original image size, we need to
# rescale it back to the testing scale to obtain RoIs.
if rescale and not isinstance(scale_factor, float):
scale_factor = torch.from_numpy(scale_factor).to(
det_bboxes.device)
if self.mask_head.bbox_type == 'hbb':
_bboxes = (
det_bboxes[:, :4] * scale_factor if rescale else det_bboxes)
else:
_bboxes = det_bboxes[:, :5]
if rescale:
_bboxes[:, :4] *= scale_factor
mask_rois = arb2roi([_bboxes], bbox_type=self.mask_head.bbox_type)
mask_results = self._mask_forward(x, mask_rois)
segm_result = self.mask_head.get_seg_masks(
mask_results['mask_pred'], _bboxes, det_labels, self.test_cfg,
ori_shape, scale_factor, rescale)
return segm_result
def simple_test(self,
x,
proposal_list,
img_metas,
proposals=None,
rescale=False):
rois = arb2roi(proposal_list, bbox_type=self.bbox_head.start_bbox_type)
bbox_results = self._bbox_forward(x, rois)
img_shape = img_metas[0]['img_shape']
scale_factor = img_metas[0]['scale_factor']
det_bboxes, det_labels = self.bbox_head.get_bboxes(
rois,
bbox_results['cls_score'],
bbox_results['bbox_pred'],
bbox_results['fix_pred'],
bbox_results['ratio_pred'],
img_shape,
scale_factor,
rescale=rescale,
cfg=self.test_cfg)
bbox_results = arb2result(det_bboxes,
det_labels,
self.bbox_head.num_classes,
bbox_type=self.bbox_head.end_bbox_type)
return bbox_results
def forward_dummy(self, x, proposals):
outs = ()
rois = arb2roi([proposals], self.bbox_head.start_bbox_type)
if self.with_bbox:
bbox_results = self._bbox_forward(x, rois)
outs = outs + (bbox_results['cls_score'],
bbox_results['bbox_pred'], bbox_results['fix_pred'],
bbox_results['ratio_pred'])
return outs
def _bbox_forward_train(self, stage, x, sampling_results, gt_bboxes,
gt_labels, rcnn_train_cfg):
rois = arb2roi([res.bboxes for res in sampling_results],
bbox_type=self.bbox_head[stage].start_bbox_type)
bbox_results = self._bbox_forward(stage, x, rois)
bbox_targets = self.bbox_head[stage].get_targets(
sampling_results, gt_bboxes, gt_labels, rcnn_train_cfg)
loss_bbox = self.bbox_head[stage].loss(bbox_results['cls_score'],
bbox_results['bbox_pred'], rois,
*bbox_targets)
bbox_results.update(
loss_bbox=loss_bbox, rois=rois, bbox_targets=bbox_targets)
return bbox_results
def forward_dummy(self, x, proposals):
"""Dummy forward function"""
# bbox head
outs = ()
rois = arb2roi([proposals], self.bbox_head.start_bbox_type)
if self.with_bbox:
bbox_results = self._bbox_forward(x, rois)
outs = outs + (bbox_results['cls_score'],
bbox_results['bbox_pred'])
# mask head
if self.with_mask:
mask_rois = rois[:100]
mask_results = self._mask_forward(x, mask_rois)
outs = outs + (mask_results['mask_pred'], )
return outs
def _bbox_forward_train(self, x, sampling_results, gt_bboxes, gt_labels,
img_metas):
"""Run forward function and calculate loss for box head in training"""
rois = arb2roi([res.bboxes for res in sampling_results],
bbox_type=self.bbox_head.start_bbox_type)
bbox_results = self._bbox_forward(x, rois)
bbox_targets = self.bbox_head.get_targets(sampling_results, gt_bboxes,
gt_labels, self.train_cfg)
loss_bbox = self.bbox_head.loss(bbox_results['cls_score'],
bbox_results['bbox_pred'], rois,
*bbox_targets)
bbox_results.update(loss_bbox=loss_bbox)
return bbox_results
def _bbox_forward_train(self, x, sampling_results, gt_bboxes, gt_labels,
img_metas):
rois = arb2roi([res.bboxes for res in sampling_results],
bbox_type=self.bbox_head.start_bbox_type)
bbox_results = self._bbox_forward(x, rois)
bbox_targets = self.bbox_head.get_targets(sampling_results, gt_bboxes,
gt_labels, self.train_cfg)
loss_bbox = self.bbox_head.loss(bbox_results['cls_score'],
bbox_results['bbox_pred'],
bbox_results['fix_pred'],
bbox_results['ratio_pred'], rois,
*bbox_targets)
bbox_results.update(loss_bbox=loss_bbox)
return bbox_results
def aug_test_bboxes(self, feats, img_metas, proposal_list, rcnn_test_cfg):
"""Test det bboxes with test time augmentation."""
aug_bboxes = []
aug_scores = []
for x, img_meta in zip(feats, img_metas):
# only one image in the batch
img_shape = img_meta[0]['img_shape']
scale_factor = img_meta[0]['scale_factor']
angle = img_meta[0].get('angle', 0)
matrix = img_meta[0].get('matrix', np.eye(3))
rotate_after_flip = img_meta[0].get('rotate_after_flip', True)
if 'flip' in img_meta[0]:
direction = img_meta[0]['flip_direction']
h_flip = img_meta[0]['flip'] if direction == 'horizontal' else False
v_flip = img_meta[0]['flip'] if direction == 'vertical' else False
else:
h_flip, v_flip = img_meta[0]['h_flip'], img_meta[0]['v_flip']
# TODO more flexible
proposals = arb_mapping(proposal_list[0][:, :-1], img_shape, scale_factor,
h_flip, v_flip, rotate_after_flip, angle, matrix,
bbox_type=self.bbox_head.start_bbox_type)
rois = arb2roi([proposals], bbox_type=self.bbox_head.start_bbox_type)
# recompute feature maps to save GPU memory
bbox_results = self._bbox_forward(x, rois)
bboxes, scores = self.bbox_head.get_bboxes(
rois,
bbox_results['cls_score'],
bbox_results['bbox_pred'],
img_shape,
scale_factor,
rescale=False,
cfg=None)
aug_bboxes.append(bboxes)
aug_scores.append(scores)
# after merging, bboxes will be rescaled to the original image size
merged_bboxes, merged_scores = merge_rotate_aug_arb(
aug_bboxes, aug_scores, img_metas, rcnn_test_cfg, merge_type='avg',
bbox_type=self.bbox_head.end_bbox_type)
det_bboxes, det_labels = multiclass_arb_nms(merged_bboxes, merged_scores,
rcnn_test_cfg.score_thr,
rcnn_test_cfg.nms,
rcnn_test_cfg.max_per_img,
bbox_type=self.bbox_head.end_bbox_type)
return det_bboxes, det_labels
def simple_test_bboxes(self,
x,
img_metas,
proposals,
rcnn_test_cfg,
rescale=False):
"""Test only det bboxes without augmentation."""
rois = arb2roi(proposals, bbox_type=self.bbox_head.start_bbox_type)
bbox_results = self._bbox_forward(x, rois)
img_shape = img_metas[0]['img_shape']
scale_factor = img_metas[0]['scale_factor']
det_bboxes, det_labels = self.bbox_head.get_bboxes(
rois,
bbox_results['cls_score'],
bbox_results['bbox_pred'],
img_shape,
scale_factor,
rescale=rescale,
cfg=rcnn_test_cfg)
return det_bboxes, det_labels
def _mask_forward_train(self, x, sampling_results, bbox_feats, gt_masks,
img_metas):
"""Run forward function and calculate loss for mask head in training"""
if not self.share_roi_extractor:
mask_head_btype = getattr(self.mask_head, 'bbox_type', 'hbb')
pos_rois = arb2roi([res.pos_bboxes for res in sampling_results],
bbox_type=mask_head_btype)
if pos_rois.shape[0] == 0:
return dict(loss_mask=None)
mask_results = self._mask_forward(x, pos_rois)
else:
pos_inds = []
device = bbox_feats.device
for res in sampling_results:
pos_inds.append(
torch.ones(
res.pos_bboxes.shape[0],
device=device,
dtype=torch.uint8))
pos_inds.append(
torch.zeros(
res.neg_bboxes.shape[0],
device=device,
dtype=torch.uint8))
pos_inds = torch.cat(pos_inds)
if pos_inds.shape[0] == 0:
return dict(loss_mask=None)
mask_results = self._mask_forward(
x, pos_inds=pos_inds, bbox_feats=bbox_feats)
mask_targets = self.mask_head.get_targets(sampling_results, gt_masks,
self.train_cfg)
pos_labels = torch.cat([res.pos_gt_labels for res in sampling_results])
loss_mask = self.mask_head.loss(mask_results['mask_pred'],
mask_targets, pos_labels)
mask_results.update(loss_mask=loss_mask, mask_targets=mask_targets)
return mask_results
async def async_test_mask(self,
x,
img_metas,
det_bboxes,
det_labels,
rescale=False,
mask_test_cfg=None):
"""Asynchronized test for mask head without augmentation."""
# image shape of the first image in the batch (only one)
ori_shape = img_metas[0]['ori_shape']
scale_factor = img_metas[0]['scale_factor']
if det_bboxes.shape[0] == 0:
segm_result = [[] for _ in range(self.mask_head.num_classes)]
else:
_bboxes = (
det_bboxes[:, :4] *
scale_factor if rescale else det_bboxes)
mask_rois = arb2roi([_bboxes], bbox_type=self.mask_head.bbox_type)
mask_feats = self.mask_roi_extractor(
x[:len(self.mask_roi_extractor.featmap_strides)],
mask_rois)
if self.with_shared_head:
mask_feats = self.shared_head(mask_feats)
if mask_test_cfg and mask_test_cfg.get('async_sleep_interval'):
sleep_interval = mask_test_cfg['async_sleep_interval']
else:
sleep_interval = 0.035
async with completed(
__name__,
'mask_head_forward',
sleep_interval=sleep_interval):
mask_pred = self.mask_head(mask_feats)
segm_result = self.mask_head.get_seg_masks(
mask_pred, _bboxes, det_labels, self.test_cfg, ori_shape,
scale_factor, rescale)
return segm_result
def simple_test(self, x, proposal_list, img_metas, rescale=False):
"""Test without augmentation."""
assert self.with_bbox, 'Bbox head must be implemented.'
img_shape = img_metas[0]['img_shape']
ori_shape = img_metas[0]['ori_shape']
scale_factor = img_metas[0]['scale_factor']
# "ms" in variable names means multi-stage
ms_scores = []
rcnn_test_cfg = self.test_cfg
rois = arb2roi(proposal_list, bbox_type=self.bbox_head[0].start_bbox_type)
for i in range(self.num_stages):
bbox_results = self._bbox_forward(i, x, rois)
ms_scores.append(bbox_results['cls_score'])
if i < self.num_stages - 1:
bbox_label = bbox_results['cls_score'].argmax(dim=1)
rois = self.bbox_head[i].regress_by_class(
rois, bbox_label, bbox_results['bbox_pred'], img_metas[0])
score_weights = rcnn_test_cfg.score_weights
assert len(score_weights) == len(ms_scores)
cls_score = sum([w * s for w, s in zip(score_weights, ms_scores)]) \
/ sum(score_weights)
det_bboxes, det_labels = self.bbox_head[-1].get_bboxes(
rois,
cls_score,
bbox_results['bbox_pred'],
img_shape,
scale_factor,
rescale=rescale,
cfg=rcnn_test_cfg)
results = arb2result(det_bboxes, det_labels, self.bbox_head[-1].num_classes,
bbox_type=self.bbox_head[-1].end_bbox_type)
return results
def aug_test(self, features, proposal_list, img_metas, rescale=False):
rcnn_test_cfg = self.test_cfg
aug_bboxes = []
aug_scores = []
for x, img_meta in zip(features, img_metas):
# only one image in the batch
img_shape = img_meta[0]['img_shape']
scale_factor = img_meta[0]['scale_factor']
angle = img_meta[0].get('angle', 0)
matrix = img_meta[0].get('matrix', np.eye(3))
rotate_after_flip = img_meta[0].get('rotate_after_flip', True)
if 'flip' in img_meta[0]:
direction = img_meta[0]['flip_direction']
h_flip = img_meta[0]['flip'] if direction == 'horizontal' else False
v_flip = img_meta[0]['flip'] if direction == 'vertical' else False
else:
h_flip, v_flip = img_meta[0]['h_flip'], img_meta[0]['v_flip']
proposals = hbb_mapping(proposal_list[0][:, :4], img_shape, scale_factor,
h_flip, v_flip, rotate_after_flip, angle, matrix)
# "ms" in variable names means multi-stage
ms_scores = []
rois = arb2roi([proposals], bbox_type=self.bbox_head[0].start_bbox_type)
for i in range(self.num_stages):
bbox_results = self._bbox_forward(i, x, rois)
ms_scores.append(bbox_results['cls_score'])
if i < self.num_stages - 1:
bbox_label = bbox_results['cls_score'].argmax(dim=1)
rois = self.bbox_head[i].regress_by_class(
rois, bbox_label, bbox_results['bbox_pred'],
img_meta[0])
score_weights = rcnn_test_cfg.score_weights
assert len(score_weights) == len(ms_scores)
cls_score = sum([w * s for w, s in zip(score_weights, ms_scores)]) \
/ sum(score_weights)
cls_score = ms_scores[0]
bboxes, scores = self.bbox_head[-1].get_bboxes(
rois,
cls_score,
bbox_results['bbox_pred'],
img_shape,
scale_factor,
rescale=False,
cfg=None)
aug_bboxes.append(bboxes)
aug_scores.append(scores)
# after merging, bboxes will be rescaled to the original image size
end_bbox_type = self.bbox_head[-1].end_bbox_type
merged_bboxes, merged_scores = merge_rotate_aug_arb(
aug_bboxes, aug_scores, img_metas, rcnn_test_cfg,
merge_type='avg', bbox_type=end_bbox_type)
det_bboxes, det_labels = multiclass_arb_nms(
merged_bboxes, merged_scores, rcnn_test_cfg.score_thr,
rcnn_test_cfg.nms, rcnn_test_cfg.max_per_img,
bbox_type=end_bbox_type)
if rescale:
_det_bboxes = det_bboxes
else:
_det_bboxes = det_bboxes.clone()
scale_factor = _det_bboxes.new_tensor(
img_metas[0][0]['scale_factor'])
if end_bbox_type == 'poly':
_det_bboxes[:, :8] *= scale_factor.repeat(2)
else:
_det_bboxes[:, :4] *= scale_factor
bbox_result = arb2result(
det_bboxes, det_labels, self.bbox_head[-1].num_classes,
bbox_type=self.bbox_head[-1].end_bbox_type)
return bbox_result