# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional, Tuple
import torch
from mmengine.structures import InstanceData
from torch import Tensor
from mmdet.registry import TASK_UTILS
from mmdet.utils import ConfigType
from .assign_result import AssignResult
from .base_assigner import BaseAssigner
def scale_boxes(bboxes: Tensor, scale: float) -> Tensor:
"""Expand an array of boxes by a given scale.
Args:
bboxes (Tensor): Shape (m, 4)
scale (float): The scale factor of bboxes
Returns:
Tensor: Shape (m, 4). Scaled bboxes
"""
assert bboxes.size(1) == 4
w_half = (bboxes[:, 2] - bboxes[:, 0]) * .5
h_half = (bboxes[:, 3] - bboxes[:, 1]) * .5
x_c = (bboxes[:, 2] + bboxes[:, 0]) * .5
y_c = (bboxes[:, 3] + bboxes[:, 1]) * .5
w_half *= scale
h_half *= scale
boxes_scaled = torch.zeros_like(bboxes)
boxes_scaled[:, 0] = x_c - w_half
boxes_scaled[:, 2] = x_c + w_half
boxes_scaled[:, 1] = y_c - h_half
boxes_scaled[:, 3] = y_c + h_half
return boxes_scaled
def is_located_in(points: Tensor, bboxes: Tensor) -> Tensor:
"""Are points located in bboxes.
Args:
points (Tensor): Points, shape: (m, 2).
bboxes (Tensor): Bounding boxes, shape: (n, 4).
Return:
Tensor: Flags indicating if points are located in bboxes,
shape: (m, n).
"""
assert points.size(1) == 2
assert bboxes.size(1) == 4
return (points[:, 0].unsqueeze(1) > bboxes[:, 0].unsqueeze(0)) & \
(points[:, 0].unsqueeze(1) < bboxes[:, 2].unsqueeze(0)) & \
(points[:, 1].unsqueeze(1) > bboxes[:, 1].unsqueeze(0)) & \
(points[:, 1].unsqueeze(1) < bboxes[:, 3].unsqueeze(0))
def bboxes_area(bboxes: Tensor) -> Tensor:
"""Compute the area of an array of bboxes.
Args:
bboxes (Tensor): The coordinates ox bboxes. Shape: (m, 4)
Returns:
Tensor: Area of the bboxes. Shape: (m, )
"""
assert bboxes.size(1) == 4
w = (bboxes[:, 2] - bboxes[:, 0])
h = (bboxes[:, 3] - bboxes[:, 1])
areas = w * h
return areas
@TASK_UTILS.register_module()
class CenterRegionAssigner(BaseAssigner):
"""Assign pixels at the center region of a bbox as positive.
Each proposals will be assigned with `-1`, `0`, or a positive integer
indicating the ground truth index.
- -1: negative samples
- semi-positive numbers: positive sample, index (0-based) of assigned gt
Args:
pos_scale (float): Threshold within which pixels are
labelled as positive.
neg_scale (float): Threshold above which pixels are
labelled as positive.
min_pos_iof (float): Minimum iof of a pixel with a gt to be
labelled as positive. Default: 1e-2
ignore_gt_scale (float): Threshold within which the pixels
are ignored when the gt is labelled as shadowed. Default: 0.5
foreground_dominate (bool): If True, the bbox will be assigned as
positive when a gt's kernel region overlaps with another's shadowed
(ignored) region, otherwise it is set as ignored. Default to False.
iou_calculator (:obj:`ConfigDict` or dict): Config of overlaps
Calculator.
"""
def __init__(
self,
pos_scale: float,
neg_scale: float,
min_pos_iof: float = 1e-2,
ignore_gt_scale: float = 0.5,
foreground_dominate: bool = False,
iou_calculator: ConfigType = dict(type='BboxOverlaps2D')
) -> None:
self.pos_scale = pos_scale
self.neg_scale = neg_scale
self.min_pos_iof = min_pos_iof
self.ignore_gt_scale = ignore_gt_scale
self.foreground_dominate = foreground_dominate
self.iou_calculator = TASK_UTILS.build(iou_calculator)
def get_gt_priorities(self, gt_bboxes: Tensor) -> Tensor:
"""Get gt priorities according to their areas.
Smaller gt has higher priority.
Args:
gt_bboxes (Tensor): Ground truth boxes, shape (k, 4).
Returns:
Tensor: The priority of gts so that gts with larger priority is
more likely to be assigned. Shape (k, )
"""
gt_areas = bboxes_area(gt_bboxes)
# Rank all gt bbox areas. Smaller objects has larger priority
_, sort_idx = gt_areas.sort(descending=True)
sort_idx = sort_idx.argsort()
return sort_idx
def assign(self,
pred_instances: InstanceData,
gt_instances: InstanceData,
gt_instances_ignore: Optional[InstanceData] = None,
**kwargs) -> AssignResult:
"""Assign gt to bboxes.
This method assigns gts to every prior (proposal/anchor), each prior
will be assigned with -1, or a semi-positive number. -1 means
negative sample, semi-positive number is the index (0-based) of
assigned gt.
Args:
pred_instances (:obj:`InstanceData`): Instances of model
predictions. It includes ``priors``, and the priors can
be anchors or points, or the bboxes predicted by the
previous stage, has shape (n, 4). The bboxes predicted by
the current model or stage will be named ``bboxes``,
``labels``, and ``scores``, the same as the ``InstanceData``
in other places.
gt_instances (:obj:`InstanceData`): Ground truth of instance
annotations. It usually includes ``bboxes``, with shape (k, 4),
and ``labels``, with shape (k, ).
gt_instances_ignore (:obj:`InstanceData`, optional): Instances
to be ignored during training. It includes ``bboxes``
attribute data that is ignored during training and testing.
Defaults to None.
Returns:
:obj:`AssignResult`: The assigned result. Note that shadowed_labels
of shape (N, 2) is also added as an `assign_result` attribute.
`shadowed_labels` is a tensor composed of N pairs of anchor_ind,
class_label], where N is the number of anchors that lie in the
outer region of a gt, anchor_ind is the shadowed anchor index
and class_label is the shadowed class label.
Example:
>>> from mmengine.structures import InstanceData
>>> self = CenterRegionAssigner(0.2, 0.2)
>>> pred_instances.priors = torch.Tensor([[0, 0, 10, 10],
... [10, 10, 20, 20]])
>>> gt_instances = InstanceData()
>>> gt_instances.bboxes = torch.Tensor([[0, 0, 10, 10]])
>>> gt_instances.labels = torch.Tensor([0])
>>> assign_result = self.assign(pred_instances, gt_instances)
>>> expected_gt_inds = torch.LongTensor([1, 0])
>>> assert torch.all(assign_result.gt_inds == expected_gt_inds)
"""
# There are in total 5 steps in the pixel assignment
# 1. Find core (the center region, say inner 0.2)
# and shadow (the relatively ourter part, say inner 0.2-0.5)
# regions of every gt.
# 2. Find all prior bboxes that lie in gt_core and gt_shadow regions
# 3. Assign prior bboxes in gt_core with a one-hot id of the gt in
# the image.
# 3.1. For overlapping objects, the prior bboxes in gt_core is
# assigned with the object with smallest area
# 4. Assign prior bboxes with class label according to its gt id.
# 4.1. Assign -1 to prior bboxes lying in shadowed gts
# 4.2. Assign positive prior boxes with the corresponding label
# 5. Find pixels lying in the shadow of an object and assign them with
# background label, but set the loss weight of its corresponding
# gt to zero.
# TODO not extract bboxes in assign.
gt_bboxes = gt_instances.bboxes
priors = pred_instances.priors
gt_labels = gt_instances.labels
assert priors.size(1) == 4, 'priors must have size of 4'
# 1. Find core positive and shadow region of every gt
gt_core = scale_boxes(gt_bboxes, self.pos_scale)
gt_shadow = scale_boxes(gt_bboxes, self.neg_scale)
# 2. Find prior bboxes that lie in gt_core and gt_shadow regions
prior_centers = (priors[:, 2:4] + priors[:, 0:2]) / 2
# The center points lie within the gt boxes
is_prior_in_gt = is_located_in(prior_centers, gt_bboxes)
# Only calculate prior and gt_core IoF. This enables small prior bboxes
# to match large gts
prior_and_gt_core_overlaps = self.iou_calculator(
priors, gt_core, mode='iof')
# The center point of effective priors should be within the gt box
is_prior_in_gt_core = is_prior_in_gt & (
prior_and_gt_core_overlaps > self.min_pos_iof) # shape (n, k)
is_prior_in_gt_shadow = (
self.iou_calculator(priors, gt_shadow, mode='iof') >
self.min_pos_iof)
# Rule out center effective positive pixels
is_prior_in_gt_shadow &= (~is_prior_in_gt_core)
num_gts, num_priors = gt_bboxes.size(0), priors.size(0)
if num_gts == 0 or num_priors == 0:
# If no gts exist, assign all pixels to negative
assigned_gt_ids = \
is_prior_in_gt_core.new_zeros((num_priors,),
dtype=torch.long)
pixels_in_gt_shadow = assigned_gt_ids.new_empty((0, 2))
else:
# Step 3: assign a one-hot gt id to each pixel, and smaller objects
# have high priority to assign the pixel.
sort_idx = self.get_gt_priorities(gt_bboxes)
assigned_gt_ids, pixels_in_gt_shadow = \
self.assign_one_hot_gt_indices(is_prior_in_gt_core,
is_prior_in_gt_shadow,
gt_priority=sort_idx)
if (gt_instances_ignore is not None
and gt_instances_ignore.bboxes.numel() > 0):
# No ground truth or boxes, return empty assignment
gt_bboxes_ignore = gt_instances_ignore.bboxes
gt_bboxes_ignore = scale_boxes(
gt_bboxes_ignore, scale=self.ignore_gt_scale)
is_prior_in_ignored_gts = is_located_in(prior_centers,
gt_bboxes_ignore)
is_prior_in_ignored_gts = is_prior_in_ignored_gts.any(dim=1)
assigned_gt_ids[is_prior_in_ignored_gts] = -1
# 4. Assign prior bboxes with class label according to its gt id.
# Default assigned label is the background (-1)
assigned_labels = assigned_gt_ids.new_full((num_priors, ), -1)
pos_inds = torch.nonzero(assigned_gt_ids > 0, as_tuple=False).squeeze()
if pos_inds.numel() > 0:
assigned_labels[pos_inds] = gt_labels[assigned_gt_ids[pos_inds] -
1]
# 5. Find pixels lying in the shadow of an object
shadowed_pixel_labels = pixels_in_gt_shadow.clone()
if pixels_in_gt_shadow.numel() > 0:
pixel_idx, gt_idx =\
pixels_in_gt_shadow[:, 0], pixels_in_gt_shadow[:, 1]
assert (assigned_gt_ids[pixel_idx] != gt_idx).all(), \
'Some pixels are dually assigned to ignore and gt!'
shadowed_pixel_labels[:, 1] = gt_labels[gt_idx - 1]
override = (
assigned_labels[pixel_idx] == shadowed_pixel_labels[:, 1])
if self.foreground_dominate:
# When a pixel is both positive and shadowed, set it as pos
shadowed_pixel_labels = shadowed_pixel_labels[~override]
else:
# When a pixel is both pos and shadowed, set it as shadowed
assigned_labels[pixel_idx[override]] = -1
assigned_gt_ids[pixel_idx[override]] = 0
assign_result = AssignResult(
num_gts, assigned_gt_ids, None, labels=assigned_labels)
# Add shadowed_labels as assign_result property. Shape: (num_shadow, 2)
assign_result.set_extra_property('shadowed_labels',
shadowed_pixel_labels)
return assign_result
def assign_one_hot_gt_indices(
self,
is_prior_in_gt_core: Tensor,
is_prior_in_gt_shadow: Tensor,
gt_priority: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
"""Assign only one gt index to each prior box.
Gts with large gt_priority are more likely to be assigned.
Args:
is_prior_in_gt_core (Tensor): Bool tensor indicating the prior
center is in the core area of a gt (e.g. 0-0.2).
Shape: (num_prior, num_gt).
is_prior_in_gt_shadow (Tensor): Bool tensor indicating the prior
center is in the shadowed area of a gt (e.g. 0.2-0.5).
Shape: (num_prior, num_gt).
gt_priority (Tensor): Priorities of gts. The gt with a higher
priority is more likely to be assigned to the bbox when the
bbox match with multiple gts. Shape: (num_gt, ).
Returns:
tuple: Returns (assigned_gt_inds, shadowed_gt_inds).
- assigned_gt_inds: The assigned gt index of each prior bbox \
(i.e. index from 1 to num_gts). Shape: (num_prior, ).
- shadowed_gt_inds: shadowed gt indices. It is a tensor of \
shape (num_ignore, 2) with first column being the shadowed prior \
bbox indices and the second column the shadowed gt \
indices (1-based).
"""
num_bboxes, num_gts = is_prior_in_gt_core.shape
if gt_priority is None:
gt_priority = torch.arange(
num_gts, device=is_prior_in_gt_core.device)
assert gt_priority.size(0) == num_gts
# The bigger gt_priority, the more preferable to be assigned
# The assigned inds are by default 0 (background)
assigned_gt_inds = is_prior_in_gt_core.new_zeros((num_bboxes, ),
dtype=torch.long)
# Shadowed bboxes are assigned to be background. But the corresponding
# label is ignored during loss calculation, which is done through
# shadowed_gt_inds
shadowed_gt_inds = torch.nonzero(is_prior_in_gt_shadow, as_tuple=False)
if is_prior_in_gt_core.sum() == 0: # No gt match
shadowed_gt_inds[:, 1] += 1 # 1-based. For consistency issue
return assigned_gt_inds, shadowed_gt_inds
# The priority of each prior box and gt pair. If one prior box is
# matched bo multiple gts. Only the pair with the highest priority
# is saved
pair_priority = is_prior_in_gt_core.new_full((num_bboxes, num_gts),
-1,
dtype=torch.long)
# Each bbox could match with multiple gts.
# The following codes deal with this situation
# Matched bboxes (to any gt). Shape: (num_pos_anchor, )
inds_of_match = torch.any(is_prior_in_gt_core, dim=1)
# The matched gt index of each positive bbox. Length >= num_pos_anchor
# , since one bbox could match multiple gts
matched_bbox_gt_inds = torch.nonzero(
is_prior_in_gt_core, as_tuple=False)[:, 1]
# Assign priority to each bbox-gt pair.
pair_priority[is_prior_in_gt_core] = gt_priority[matched_bbox_gt_inds]
_, argmax_priority = pair_priority[inds_of_match].max(dim=1)
assigned_gt_inds[inds_of_match] = argmax_priority + 1 # 1-based
# Zero-out the assigned anchor box to filter the shadowed gt indices
is_prior_in_gt_core[inds_of_match, argmax_priority] = 0
# Concat the shadowed indices due to overlapping with that out side of
# effective scale. shape: (total_num_ignore, 2)
shadowed_gt_inds = torch.cat(
(shadowed_gt_inds,
torch.nonzero(is_prior_in_gt_core, as_tuple=False)),
dim=0)
# Change `is_prior_in_gt_core` back to keep arguments intact.
is_prior_in_gt_core[inds_of_match, argmax_priority] = 1
# 1-based shadowed gt indices, to be consistent with `assigned_gt_inds`
if shadowed_gt_inds.numel() > 0:
shadowed_gt_inds[:, 1] += 1
return assigned_gt_inds, shadowed_gt_inds