Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from typing import Optional, Sequence, Tuple, Union | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import Tensor | |
| from mmdet.registry import TASK_UTILS | |
| from mmdet.structures.bbox import (BaseBoxes, HorizontalBoxes, bbox_rescale, | |
| get_box_tensor) | |
| from .base_bbox_coder import BaseBBoxCoder | |
| class BucketingBBoxCoder(BaseBBoxCoder): | |
| """Bucketing BBox Coder for Side-Aware Boundary Localization (SABL). | |
| Boundary Localization with Bucketing and Bucketing Guided Rescoring | |
| are implemented here. | |
| Please refer to https://arxiv.org/abs/1912.04260 for more details. | |
| Args: | |
| num_buckets (int): Number of buckets. | |
| scale_factor (int): Scale factor of proposals to generate buckets. | |
| offset_topk (int): Topk buckets are used to generate | |
| bucket fine regression targets. Defaults to 2. | |
| offset_upperbound (float): Offset upperbound to generate | |
| bucket fine regression targets. | |
| To avoid too large offset displacements. Defaults to 1.0. | |
| cls_ignore_neighbor (bool): Ignore second nearest bucket or Not. | |
| Defaults to True. | |
| clip_border (bool, optional): Whether clip the objects outside the | |
| border of the image. Defaults to True. | |
| """ | |
| def __init__(self, | |
| num_buckets: int, | |
| scale_factor: int, | |
| offset_topk: int = 2, | |
| offset_upperbound: float = 1.0, | |
| cls_ignore_neighbor: bool = True, | |
| clip_border: bool = True, | |
| **kwargs) -> None: | |
| super().__init__(**kwargs) | |
| self.num_buckets = num_buckets | |
| self.scale_factor = scale_factor | |
| self.offset_topk = offset_topk | |
| self.offset_upperbound = offset_upperbound | |
| self.cls_ignore_neighbor = cls_ignore_neighbor | |
| self.clip_border = clip_border | |
| def encode(self, bboxes: Union[Tensor, BaseBoxes], | |
| gt_bboxes: Union[Tensor, BaseBoxes]) -> Tuple[Tensor]: | |
| """Get bucketing estimation and fine regression targets during | |
| training. | |
| Args: | |
| bboxes (torch.Tensor or :obj:`BaseBoxes`): source boxes, | |
| e.g., object proposals. | |
| gt_bboxes (torch.Tensor or :obj:`BaseBoxes`): target of the | |
| transformation, e.g., ground truth boxes. | |
| Returns: | |
| encoded_bboxes(tuple[Tensor]): bucketing estimation | |
| and fine regression targets and weights | |
| """ | |
| bboxes = get_box_tensor(bboxes) | |
| gt_bboxes = get_box_tensor(gt_bboxes) | |
| assert bboxes.size(0) == gt_bboxes.size(0) | |
| assert bboxes.size(-1) == gt_bboxes.size(-1) == 4 | |
| encoded_bboxes = bbox2bucket(bboxes, gt_bboxes, self.num_buckets, | |
| self.scale_factor, self.offset_topk, | |
| self.offset_upperbound, | |
| self.cls_ignore_neighbor) | |
| return encoded_bboxes | |
| def decode( | |
| self, | |
| bboxes: Union[Tensor, BaseBoxes], | |
| pred_bboxes: Tensor, | |
| max_shape: Optional[Tuple[int]] = None | |
| ) -> Tuple[Union[Tensor, BaseBoxes], Tensor]: | |
| """Apply transformation `pred_bboxes` to `boxes`. | |
| Args: | |
| boxes (torch.Tensor or :obj:`BaseBoxes`): Basic boxes. | |
| pred_bboxes (torch.Tensor): Predictions for bucketing estimation | |
| and fine regression | |
| max_shape (tuple[int], optional): Maximum shape of boxes. | |
| Defaults to None. | |
| Returns: | |
| Union[torch.Tensor, :obj:`BaseBoxes`]: Decoded boxes. | |
| """ | |
| bboxes = get_box_tensor(bboxes) | |
| assert len(pred_bboxes) == 2 | |
| cls_preds, offset_preds = pred_bboxes | |
| assert cls_preds.size(0) == bboxes.size(0) and offset_preds.size( | |
| 0) == bboxes.size(0) | |
| bboxes, loc_confidence = bucket2bbox(bboxes, cls_preds, offset_preds, | |
| self.num_buckets, | |
| self.scale_factor, max_shape, | |
| self.clip_border) | |
| if self.use_box_type: | |
| bboxes = HorizontalBoxes(bboxes, clone=False) | |
| return bboxes, loc_confidence | |
| def generat_buckets(proposals: Tensor, | |
| num_buckets: int, | |
| scale_factor: float = 1.0) -> Tuple[Tensor]: | |
| """Generate buckets w.r.t bucket number and scale factor of proposals. | |
| Args: | |
| proposals (Tensor): Shape (n, 4) | |
| num_buckets (int): Number of buckets. | |
| scale_factor (float): Scale factor to rescale proposals. | |
| Returns: | |
| tuple[Tensor]: (bucket_w, bucket_h, l_buckets, r_buckets, | |
| t_buckets, d_buckets) | |
| - bucket_w: Width of buckets on x-axis. Shape (n, ). | |
| - bucket_h: Height of buckets on y-axis. Shape (n, ). | |
| - l_buckets: Left buckets. Shape (n, ceil(side_num/2)). | |
| - r_buckets: Right buckets. Shape (n, ceil(side_num/2)). | |
| - t_buckets: Top buckets. Shape (n, ceil(side_num/2)). | |
| - d_buckets: Down buckets. Shape (n, ceil(side_num/2)). | |
| """ | |
| proposals = bbox_rescale(proposals, scale_factor) | |
| # number of buckets in each side | |
| side_num = int(np.ceil(num_buckets / 2.0)) | |
| pw = proposals[..., 2] - proposals[..., 0] | |
| ph = proposals[..., 3] - proposals[..., 1] | |
| px1 = proposals[..., 0] | |
| py1 = proposals[..., 1] | |
| px2 = proposals[..., 2] | |
| py2 = proposals[..., 3] | |
| bucket_w = pw / num_buckets | |
| bucket_h = ph / num_buckets | |
| # left buckets | |
| l_buckets = px1[:, None] + (0.5 + torch.arange( | |
| 0, side_num).to(proposals).float())[None, :] * bucket_w[:, None] | |
| # right buckets | |
| r_buckets = px2[:, None] - (0.5 + torch.arange( | |
| 0, side_num).to(proposals).float())[None, :] * bucket_w[:, None] | |
| # top buckets | |
| t_buckets = py1[:, None] + (0.5 + torch.arange( | |
| 0, side_num).to(proposals).float())[None, :] * bucket_h[:, None] | |
| # down buckets | |
| d_buckets = py2[:, None] - (0.5 + torch.arange( | |
| 0, side_num).to(proposals).float())[None, :] * bucket_h[:, None] | |
| return bucket_w, bucket_h, l_buckets, r_buckets, t_buckets, d_buckets | |
| def bbox2bucket(proposals: Tensor, | |
| gt: Tensor, | |
| num_buckets: int, | |
| scale_factor: float, | |
| offset_topk: int = 2, | |
| offset_upperbound: float = 1.0, | |
| cls_ignore_neighbor: bool = True) -> Tuple[Tensor]: | |
| """Generate buckets estimation and fine regression targets. | |
| Args: | |
| proposals (Tensor): Shape (n, 4) | |
| gt (Tensor): Shape (n, 4) | |
| num_buckets (int): Number of buckets. | |
| scale_factor (float): Scale factor to rescale proposals. | |
| offset_topk (int): Topk buckets are used to generate | |
| bucket fine regression targets. Defaults to 2. | |
| offset_upperbound (float): Offset allowance to generate | |
| bucket fine regression targets. | |
| To avoid too large offset displacements. Defaults to 1.0. | |
| cls_ignore_neighbor (bool): Ignore second nearest bucket or Not. | |
| Defaults to True. | |
| Returns: | |
| tuple[Tensor]: (offsets, offsets_weights, bucket_labels, cls_weights). | |
| - offsets: Fine regression targets. \ | |
| Shape (n, num_buckets*2). | |
| - offsets_weights: Fine regression weights. \ | |
| Shape (n, num_buckets*2). | |
| - bucket_labels: Bucketing estimation labels. \ | |
| Shape (n, num_buckets*2). | |
| - cls_weights: Bucketing estimation weights. \ | |
| Shape (n, num_buckets*2). | |
| """ | |
| assert proposals.size() == gt.size() | |
| # generate buckets | |
| proposals = proposals.float() | |
| gt = gt.float() | |
| (bucket_w, bucket_h, l_buckets, r_buckets, t_buckets, | |
| d_buckets) = generat_buckets(proposals, num_buckets, scale_factor) | |
| gx1 = gt[..., 0] | |
| gy1 = gt[..., 1] | |
| gx2 = gt[..., 2] | |
| gy2 = gt[..., 3] | |
| # generate offset targets and weights | |
| # offsets from buckets to gts | |
| l_offsets = (l_buckets - gx1[:, None]) / bucket_w[:, None] | |
| r_offsets = (r_buckets - gx2[:, None]) / bucket_w[:, None] | |
| t_offsets = (t_buckets - gy1[:, None]) / bucket_h[:, None] | |
| d_offsets = (d_buckets - gy2[:, None]) / bucket_h[:, None] | |
| # select top-k nearest buckets | |
| l_topk, l_label = l_offsets.abs().topk( | |
| offset_topk, dim=1, largest=False, sorted=True) | |
| r_topk, r_label = r_offsets.abs().topk( | |
| offset_topk, dim=1, largest=False, sorted=True) | |
| t_topk, t_label = t_offsets.abs().topk( | |
| offset_topk, dim=1, largest=False, sorted=True) | |
| d_topk, d_label = d_offsets.abs().topk( | |
| offset_topk, dim=1, largest=False, sorted=True) | |
| offset_l_weights = l_offsets.new_zeros(l_offsets.size()) | |
| offset_r_weights = r_offsets.new_zeros(r_offsets.size()) | |
| offset_t_weights = t_offsets.new_zeros(t_offsets.size()) | |
| offset_d_weights = d_offsets.new_zeros(d_offsets.size()) | |
| inds = torch.arange(0, proposals.size(0)).to(proposals).long() | |
| # generate offset weights of top-k nearest buckets | |
| for k in range(offset_topk): | |
| if k >= 1: | |
| offset_l_weights[inds, l_label[:, | |
| k]] = (l_topk[:, k] < | |
| offset_upperbound).float() | |
| offset_r_weights[inds, r_label[:, | |
| k]] = (r_topk[:, k] < | |
| offset_upperbound).float() | |
| offset_t_weights[inds, t_label[:, | |
| k]] = (t_topk[:, k] < | |
| offset_upperbound).float() | |
| offset_d_weights[inds, d_label[:, | |
| k]] = (d_topk[:, k] < | |
| offset_upperbound).float() | |
| else: | |
| offset_l_weights[inds, l_label[:, k]] = 1.0 | |
| offset_r_weights[inds, r_label[:, k]] = 1.0 | |
| offset_t_weights[inds, t_label[:, k]] = 1.0 | |
| offset_d_weights[inds, d_label[:, k]] = 1.0 | |
| offsets = torch.cat([l_offsets, r_offsets, t_offsets, d_offsets], dim=-1) | |
| offsets_weights = torch.cat([ | |
| offset_l_weights, offset_r_weights, offset_t_weights, offset_d_weights | |
| ], | |
| dim=-1) | |
| # generate bucket labels and weight | |
| side_num = int(np.ceil(num_buckets / 2.0)) | |
| labels = torch.stack( | |
| [l_label[:, 0], r_label[:, 0], t_label[:, 0], d_label[:, 0]], dim=-1) | |
| batch_size = labels.size(0) | |
| bucket_labels = F.one_hot(labels.view(-1), side_num).view(batch_size, | |
| -1).float() | |
| bucket_cls_l_weights = (l_offsets.abs() < 1).float() | |
| bucket_cls_r_weights = (r_offsets.abs() < 1).float() | |
| bucket_cls_t_weights = (t_offsets.abs() < 1).float() | |
| bucket_cls_d_weights = (d_offsets.abs() < 1).float() | |
| bucket_cls_weights = torch.cat([ | |
| bucket_cls_l_weights, bucket_cls_r_weights, bucket_cls_t_weights, | |
| bucket_cls_d_weights | |
| ], | |
| dim=-1) | |
| # ignore second nearest buckets for cls if necessary | |
| if cls_ignore_neighbor: | |
| bucket_cls_weights = (~((bucket_cls_weights == 1) & | |
| (bucket_labels == 0))).float() | |
| else: | |
| bucket_cls_weights[:] = 1.0 | |
| return offsets, offsets_weights, bucket_labels, bucket_cls_weights | |
| def bucket2bbox(proposals: Tensor, | |
| cls_preds: Tensor, | |
| offset_preds: Tensor, | |
| num_buckets: int, | |
| scale_factor: float = 1.0, | |
| max_shape: Optional[Union[Sequence[int], Tensor, | |
| Sequence[Sequence[int]]]] = None, | |
| clip_border: bool = True) -> Tuple[Tensor]: | |
| """Apply bucketing estimation (cls preds) and fine regression (offset | |
| preds) to generate det bboxes. | |
| Args: | |
| proposals (Tensor): Boxes to be transformed. Shape (n, 4) | |
| cls_preds (Tensor): bucketing estimation. Shape (n, num_buckets*2). | |
| offset_preds (Tensor): fine regression. Shape (n, num_buckets*2). | |
| num_buckets (int): Number of buckets. | |
| scale_factor (float): Scale factor to rescale proposals. | |
| max_shape (tuple[int, int]): Maximum bounds for boxes. specifies (H, W) | |
| clip_border (bool, optional): Whether clip the objects outside the | |
| border of the image. Defaults to True. | |
| Returns: | |
| tuple[Tensor]: (bboxes, loc_confidence). | |
| - bboxes: predicted bboxes. Shape (n, 4) | |
| - loc_confidence: localization confidence of predicted bboxes. | |
| Shape (n,). | |
| """ | |
| side_num = int(np.ceil(num_buckets / 2.0)) | |
| cls_preds = cls_preds.view(-1, side_num) | |
| offset_preds = offset_preds.view(-1, side_num) | |
| scores = F.softmax(cls_preds, dim=1) | |
| score_topk, score_label = scores.topk(2, dim=1, largest=True, sorted=True) | |
| rescaled_proposals = bbox_rescale(proposals, scale_factor) | |
| pw = rescaled_proposals[..., 2] - rescaled_proposals[..., 0] | |
| ph = rescaled_proposals[..., 3] - rescaled_proposals[..., 1] | |
| px1 = rescaled_proposals[..., 0] | |
| py1 = rescaled_proposals[..., 1] | |
| px2 = rescaled_proposals[..., 2] | |
| py2 = rescaled_proposals[..., 3] | |
| bucket_w = pw / num_buckets | |
| bucket_h = ph / num_buckets | |
| score_inds_l = score_label[0::4, 0] | |
| score_inds_r = score_label[1::4, 0] | |
| score_inds_t = score_label[2::4, 0] | |
| score_inds_d = score_label[3::4, 0] | |
| l_buckets = px1 + (0.5 + score_inds_l.float()) * bucket_w | |
| r_buckets = px2 - (0.5 + score_inds_r.float()) * bucket_w | |
| t_buckets = py1 + (0.5 + score_inds_t.float()) * bucket_h | |
| d_buckets = py2 - (0.5 + score_inds_d.float()) * bucket_h | |
| offsets = offset_preds.view(-1, 4, side_num) | |
| inds = torch.arange(proposals.size(0)).to(proposals).long() | |
| l_offsets = offsets[:, 0, :][inds, score_inds_l] | |
| r_offsets = offsets[:, 1, :][inds, score_inds_r] | |
| t_offsets = offsets[:, 2, :][inds, score_inds_t] | |
| d_offsets = offsets[:, 3, :][inds, score_inds_d] | |
| x1 = l_buckets - l_offsets * bucket_w | |
| x2 = r_buckets - r_offsets * bucket_w | |
| y1 = t_buckets - t_offsets * bucket_h | |
| y2 = d_buckets - d_offsets * bucket_h | |
| if clip_border and max_shape is not None: | |
| x1 = x1.clamp(min=0, max=max_shape[1] - 1) | |
| y1 = y1.clamp(min=0, max=max_shape[0] - 1) | |
| x2 = x2.clamp(min=0, max=max_shape[1] - 1) | |
| y2 = y2.clamp(min=0, max=max_shape[0] - 1) | |
| bboxes = torch.cat([x1[:, None], y1[:, None], x2[:, None], y2[:, None]], | |
| dim=-1) | |
| # bucketing guided rescoring | |
| loc_confidence = score_topk[:, 0] | |
| top2_neighbor_inds = (score_label[:, 0] - score_label[:, 1]).abs() == 1 | |
| loc_confidence += score_topk[:, 1] * top2_neighbor_inds.float() | |
| loc_confidence = loc_confidence.view(-1, 4).mean(dim=1) | |
| return bboxes, loc_confidence | |