227 lines
9.9 KiB
Python
227 lines
9.9 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import numpy as np
|
|
import paddle
|
|
import paddle.nn as nn
|
|
import paddle.nn.functional as F
|
|
|
|
from ppdet.core.workspace import register
|
|
from ppdet.modeling.rbox_utils import box2corners, check_points_in_polys, paddle_gather
|
|
|
|
__all__ = ['FCOSRAssigner']
|
|
|
|
EPS = 1e-9
|
|
|
|
|
|
@register
|
|
class FCOSRAssigner(nn.Layer):
|
|
""" FCOSR Assigner, refer to https://arxiv.org/abs/2111.10780 for details
|
|
|
|
1. compute normalized gaussian distribution score and refined gaussian distribution score
|
|
2. refer to ellipse center sampling, sample points whose normalized gaussian distribution score is greater than threshold
|
|
3. refer to multi-level sampling, assign ground truth to feature map which follows two conditions.
|
|
i). first, the ratio between the short edge of the target and the stride of the feature map is less than 2.
|
|
ii). second, the long edge of minimum bounding rectangle of the target is larger than the acceptance range of feature map
|
|
4. refer to fuzzy sample label assignment, the points satisfying 2 and 3 will be assigned to the ground truth according to gaussian distribution score
|
|
"""
|
|
__shared__ = ['num_classes']
|
|
|
|
def __init__(self,
|
|
num_classes=80,
|
|
factor=12,
|
|
threshold=0.23,
|
|
boundary=[[-1, 128], [128, 320], [320, 10000]],
|
|
score_type='iou'):
|
|
super(FCOSRAssigner, self).__init__()
|
|
self.num_classes = num_classes
|
|
self.factor = factor
|
|
self.threshold = threshold
|
|
self.boundary = [
|
|
paddle.to_tensor(
|
|
l, dtype=paddle.float32).reshape([1, 1, 2]) for l in boundary
|
|
]
|
|
self.score_type = score_type
|
|
|
|
def get_gaussian_distribution_score(self, points, gt_rboxes, gt_polys):
|
|
# projecting points to coordinate system defined by each rbox
|
|
# [B, N, 4, 2] -> 4 * [B, N, 1, 2]
|
|
a, b, c, d = gt_polys.split(4, axis=2)
|
|
# [1, L, 2] -> [1, 1, L, 2]
|
|
points = points.unsqueeze(0)
|
|
ab = b - a
|
|
ad = d - a
|
|
# [B, N, 5] -> [B, N, 2], [B, N, 2], [B, N, 1]
|
|
xy, wh, angle = gt_rboxes.split([2, 2, 1], axis=-1)
|
|
# [B, N, 2] -> [B, N, 1, 2]
|
|
xy = xy.unsqueeze(2)
|
|
# vector of points to center [B, N, L, 2]
|
|
vec = points - xy
|
|
# <ab, vec> = |ab| * |vec| * cos(theta) [B, N, L]
|
|
vec_dot_ab = paddle.sum(vec * ab, axis=-1)
|
|
# <ad, vec> = |ad| * |vec| * cos(theta) [B, N, L]
|
|
vec_dot_ad = paddle.sum(vec * ad, axis=-1)
|
|
# norm_ab [B, N, L]
|
|
norm_ab = paddle.sum(ab * ab, axis=-1).sqrt()
|
|
# norm_ad [B, N, L]
|
|
norm_ad = paddle.sum(ad * ad, axis=-1).sqrt()
|
|
# min(h, w), [B, N, 1]
|
|
min_edge = paddle.min(wh, axis=-1, keepdim=True)
|
|
# delta_x, delta_y [B, N, L]
|
|
delta_x = vec_dot_ab.pow(2) / (norm_ab.pow(3) * min_edge + EPS)
|
|
delta_y = vec_dot_ad.pow(2) / (norm_ad.pow(3) * min_edge + EPS)
|
|
# score [B, N, L]
|
|
norm_score = paddle.exp(-0.5 * self.factor * (delta_x + delta_y))
|
|
|
|
# simplified calculation
|
|
sigma = min_edge / self.factor
|
|
refined_score = norm_score / (2 * np.pi * sigma + EPS)
|
|
return norm_score, refined_score
|
|
|
|
def get_rotated_inside_mask(self, points, gt_polys, scores):
|
|
inside_mask = check_points_in_polys(points, gt_polys)
|
|
center_mask = scores >= self.threshold
|
|
return (inside_mask & center_mask).cast(paddle.float32)
|
|
|
|
def get_inside_range_mask(self, points, gt_bboxes, gt_rboxes, stride_tensor,
|
|
regress_range):
|
|
# [1, L, 2] -> [1, 1, L, 2]
|
|
points = points.unsqueeze(0)
|
|
# [B, n, 4] -> [B, n, 1, 4]
|
|
x1y1, x2y2 = gt_bboxes.unsqueeze(2).split(2, axis=-1)
|
|
# [B, n, L, 2]
|
|
lt = points - x1y1
|
|
rb = x2y2 - points
|
|
# [B, n, L, 4]
|
|
ltrb = paddle.concat([lt, rb], axis=-1)
|
|
# [B, n, L, 4] -> [B, n, L]
|
|
inside_mask = paddle.min(ltrb, axis=-1) > EPS
|
|
# regress_range [1, L, 2] -> [1, 1, L, 2]
|
|
regress_range = regress_range.unsqueeze(0)
|
|
# stride_tensor [1, L, 1] -> [1, 1, L]
|
|
stride_tensor = stride_tensor.transpose((0, 2, 1))
|
|
# fcos range
|
|
# [B, n, L, 4] -> [B, n, L]
|
|
ltrb_max = paddle.max(ltrb, axis=-1)
|
|
# [1, 1, L, 2] -> [1, 1, L]
|
|
low, high = regress_range[..., 0], regress_range[..., 1]
|
|
# [B, n, L]
|
|
regress_mask = (ltrb_max >= low) & (ltrb_max <= high)
|
|
# mask for rotated
|
|
# [B, n, 1]
|
|
min_edge = paddle.min(gt_rboxes[..., 2:4], axis=-1, keepdim=True)
|
|
# [B, n , L]
|
|
rotated_mask = ((min_edge / stride_tensor) < 2.0) & (ltrb_max > high)
|
|
mask = inside_mask & (regress_mask | rotated_mask)
|
|
return mask.cast(paddle.float32)
|
|
|
|
@paddle.no_grad()
|
|
def forward(self,
|
|
anchor_points,
|
|
stride_tensor,
|
|
num_anchors_list,
|
|
gt_labels,
|
|
gt_bboxes,
|
|
gt_rboxes,
|
|
pad_gt_mask,
|
|
bg_index,
|
|
pred_rboxes=None):
|
|
r"""
|
|
|
|
Args:
|
|
anchor_points (Tensor, float32): pre-defined anchor points, shape(1, L, 2),
|
|
"x, y" format
|
|
stride_tensor (Tensor, float32): stride tensor, shape (1, L, 1)
|
|
num_anchors_list (List): num of anchors in each level
|
|
gt_labels (Tensor, int64|int32): Label of gt_bboxes, shape(B, n, 1)
|
|
gt_bboxes (Tensor, float32): Ground truth bboxes, shape(B, n, 4)
|
|
gt_rboxes (Tensor, float32): Ground truth bboxes, shape(B, n, 5)
|
|
pad_gt_mask (Tensor, float32): 1 means bbox, 0 means no bbox, shape(B, n, 1)
|
|
bg_index (int): background index
|
|
pred_rboxes (Tensor, float32, optional): predicted bounding boxes, shape(B, L, 5)
|
|
Returns:
|
|
assigned_labels (Tensor): (B, L)
|
|
assigned_rboxes (Tensor): (B, L, 5)
|
|
assigned_scores (Tensor): (B, L, C), if pred_rboxes is not None, then output ious
|
|
"""
|
|
|
|
_, num_anchors, _ = anchor_points.shape
|
|
batch_size, num_max_boxes, _ = gt_rboxes.shape
|
|
if num_max_boxes == 0:
|
|
assigned_labels = paddle.full(
|
|
[batch_size, num_anchors], bg_index, dtype=gt_labels.dtype)
|
|
assigned_rboxes = paddle.zeros([batch_size, num_anchors, 5])
|
|
assigned_scores = paddle.zeros(
|
|
[batch_size, num_anchors, self.num_classes])
|
|
return assigned_labels, assigned_rboxes, assigned_scores
|
|
|
|
# get normalized gaussian distribution score and refined distribution score
|
|
gt_polys = box2corners(gt_rboxes)
|
|
score, refined_score = self.get_gaussian_distribution_score(
|
|
anchor_points, gt_rboxes, gt_polys)
|
|
inside_mask = self.get_rotated_inside_mask(anchor_points, gt_polys,
|
|
score)
|
|
regress_ranges = []
|
|
for num, bound in zip(num_anchors_list, self.boundary):
|
|
regress_ranges.append(bound.tile((1, num, 1)))
|
|
regress_ranges = paddle.concat(regress_ranges, axis=1)
|
|
regress_mask = self.get_inside_range_mask(
|
|
anchor_points, gt_bboxes, gt_rboxes, stride_tensor, regress_ranges)
|
|
# [B, n, L]
|
|
mask_positive = inside_mask * regress_mask * pad_gt_mask
|
|
refined_score = refined_score * mask_positive - (1. - mask_positive)
|
|
|
|
argmax_refined_score = refined_score.argmax(axis=-2)
|
|
max_refined_score = refined_score.max(axis=-2)
|
|
assigned_gt_index = argmax_refined_score
|
|
|
|
# assigned target
|
|
batch_ind = paddle.arange(
|
|
end=batch_size, dtype=gt_labels.dtype).unsqueeze(-1)
|
|
assigned_gt_index = assigned_gt_index + batch_ind * num_max_boxes
|
|
assigned_labels = paddle.gather(
|
|
gt_labels.flatten(), assigned_gt_index.flatten(), axis=0)
|
|
assigned_labels = assigned_labels.reshape([batch_size, num_anchors])
|
|
assigned_labels = paddle.where(
|
|
max_refined_score > 0, assigned_labels,
|
|
paddle.full_like(assigned_labels, bg_index))
|
|
|
|
assigned_rboxes = paddle.gather(
|
|
gt_rboxes.reshape([-1, 5]), assigned_gt_index.flatten(), axis=0)
|
|
assigned_rboxes = assigned_rboxes.reshape([batch_size, num_anchors, 5])
|
|
|
|
assigned_scores = F.one_hot(assigned_labels, self.num_classes + 1)
|
|
ind = list(range(self.num_classes + 1))
|
|
ind.remove(bg_index)
|
|
assigned_scores = paddle.index_select(
|
|
assigned_scores, paddle.to_tensor(ind), axis=-1)
|
|
|
|
if self.score_type == 'gaussian':
|
|
selected_scores = paddle_gather(
|
|
score, 1, argmax_refined_score.unsqueeze(-2)).squeeze(-2)
|
|
assigned_scores = assigned_scores * selected_scores.unsqueeze(-1)
|
|
elif self.score_type == 'iou':
|
|
assert pred_rboxes is not None, 'If score type is iou, pred_rboxes should not be None'
|
|
from ext_op import matched_rbox_iou
|
|
b, l = pred_rboxes.shape[:2]
|
|
iou_score = matched_rbox_iou(
|
|
pred_rboxes.reshape((-1, 5)), assigned_rboxes.reshape(
|
|
(-1, 5))).reshape((b, l, 1))
|
|
assigned_scores = assigned_scores * iou_score
|
|
|
|
return assigned_labels, assigned_rboxes, assigned_scores |