#!/usr/bin/env python3
# Copyright (c) OpenMMLab. All rights reserved.
# Copyright (c) Megvii Inc. All rights reserved.
# Copyright (C) Alibaba Group Holding Limited. All rights reserved.
import numpy as np
import torch
import torchvision
from ..structures.bounding_box import BoxList
__all__ = [
'filter_box',
'postprocess',
'bboxes_iou',
'matrix_iou',
'adjust_box_anns',
'xyxy2xywh',
'xyxy2cxcywh',
]
def multiclass_nms(multi_bboxes,
multi_scores,
score_thr,
iou_thr,
max_num=100,
score_factors=None):
"""NMS for multi-class bboxes.
Args:
multi_bboxes (Tensor): shape (n, #class*4) or (n, 4)
multi_scores (Tensor): shape (n, #class), where the last column
contains scores of the background class, but this will be ignored.
score_thr (float): bbox threshold, bboxes with scores lower than it
will not be considered.
nms_thr (float): NMS IoU threshold
max_num (int): if there are more than max_num bboxes after NMS,
only top max_num will be kept.
score_factors (Tensor): The factors multiplied to scores before
applying NMS
Returns:
tuple: (bboxes, labels), tensors of shape (k, 5) and (k, 1). Labels \
are 0-based.
"""
num_classes = multi_scores.size(1)
# exclude background category
if multi_bboxes.shape[1] > 4:
bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
else:
bboxes = multi_bboxes[:, None].expand(multi_scores.size(0),
num_classes, 4)
scores = multi_scores
# filter out boxes with low scores
valid_mask = scores > score_thr # 1000 * 80 bool
# We use masked_select for ONNX exporting purpose,
# which is equivalent to bboxes = bboxes[valid_mask]
# (TODO): as ONNX does not support repeat now,
# we have to use this ugly code
# bboxes -> 1000, 4
bboxes = torch.masked_select(
bboxes,
torch.stack((valid_mask, valid_mask, valid_mask, valid_mask),
-1)).view(-1, 4) # mask-> 1000*80*4, 80000*4
if score_factors is not None:
scores = scores * score_factors[:, None]
scores = torch.masked_select(scores, valid_mask)
labels = valid_mask.nonzero(as_tuple=False)[:, 1]
if bboxes.numel() == 0:
bboxes = multi_bboxes.new_zeros((0, 5))
labels = multi_bboxes.new_zeros((0, ), dtype=torch.long)
scores = multi_bboxes.new_zeros((0, ))
return bboxes, scores, labels
keep = torchvision.ops.batched_nms(bboxes, scores, labels, iou_thr)
if max_num > 0:
keep = keep[:max_num]
return bboxes[keep], scores[keep], labels[keep]
[ドキュメント]
def filter_box(output, scale_range):
"""
output: (N, 5+class) shape
"""
min_scale, max_scale = scale_range
w = output[:, 2] - output[:, 0]
h = output[:, 3] - output[:, 1]
keep = (w * h > min_scale * min_scale) & (w * h < max_scale * max_scale)
return output[keep]
def filter_results(boxlist, num_classes, nms_thre):
boxes = boxlist.bbox
scores = boxlist.get_field('scores')
cls = boxlist.get_field('labels')
nms_out_index = torchvision.ops.batched_nms(
boxes,
scores,
cls,
nms_thre,
)
boxlist = boxlist[nms_out_index]
return boxlist
[ドキュメント]
def postprocess(cls_scores,
bbox_preds,
num_classes,
conf_thre=0.7,
nms_thre=0.45,
imgs=None):
batch_size = bbox_preds.size(0)
output = [None for _ in range(batch_size)]
for i in range(batch_size):
# If none are remaining => process next image
if not bbox_preds[i].size(0):
continue
detections, scores, labels = multiclass_nms(bbox_preds[i],
cls_scores[i], conf_thre,
nms_thre, 500)
detections = torch.cat((detections, torch.ones_like(
scores[:, None]), scores[:, None], labels[:, None]),
dim=1)
if output[i] is None:
output[i] = detections
else:
output[i] = torch.cat((output[i], detections))
# transfer to BoxList
for i in range(len(output)):
res = output[i]
if res is None or imgs is None:
boxlist = BoxList(torch.zeros(0, 4), (0, 0), mode='xyxy')
boxlist.add_field('objectness', 0)
boxlist.add_field('scores', 0)
boxlist.add_field('labels', -1)
else:
img_h, img_w = imgs.image_sizes[i]
boxlist = BoxList(res[:, :4], (img_w, img_h), mode='xyxy')
boxlist.add_field('objectness', res[:, 4])
boxlist.add_field('scores', res[:, 5])
boxlist.add_field('labels', res[:, 6])
output[i] = boxlist
return output
[ドキュメント]
def bboxes_iou(bboxes_a, bboxes_b, xyxy=True):
if bboxes_a.shape[1] != 4 or bboxes_b.shape[1] != 4:
raise IndexError
if xyxy:
tl = torch.max(bboxes_a[:, None, :2], bboxes_b[:, :2])
br = torch.min(bboxes_a[:, None, 2:], bboxes_b[:, 2:])
area_a = torch.prod(bboxes_a[:, 2:] - bboxes_a[:, :2], 1)
area_b = torch.prod(bboxes_b[:, 2:] - bboxes_b[:, :2], 1)
else:
tl = torch.max(
(bboxes_a[:, None, :2] - bboxes_a[:, None, 2:] / 2),
(bboxes_b[:, :2] - bboxes_b[:, 2:] / 2),
)
br = torch.min(
(bboxes_a[:, None, :2] + bboxes_a[:, None, 2:] / 2),
(bboxes_b[:, :2] + bboxes_b[:, 2:] / 2),
)
area_a = torch.prod(bboxes_a[:, 2:], 1)
area_b = torch.prod(bboxes_b[:, 2:], 1)
en = (tl < br).type(tl.type()).prod(dim=2)
area_i = torch.prod(br - tl, 2) * en # * ((tl < br).all())
return area_i / (area_a[:, None] + area_b - area_i)
[ドキュメント]
def matrix_iou(a, b):
"""
return iou of a and b, numpy version for data augenmentation
"""
lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
return area_i / (area_a[:, np.newaxis] + area_b - area_i + 1e-12)
[ドキュメント]
def adjust_box_anns(bbox, scale_ratio, padw, padh, w_max, h_max):
bbox[:, 0::2] = np.clip(bbox[:, 0::2] * scale_ratio + padw, 0, w_max)
bbox[:, 1::2] = np.clip(bbox[:, 1::2] * scale_ratio + padh, 0, h_max)
return bbox
[ドキュメント]
def xyxy2xywh(bboxes):
bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
return bboxes
[ドキュメント]
def xyxy2cxcywh(bboxes):
bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
bboxes[:, 0] = bboxes[:, 0] + bboxes[:, 2] * 0.5
bboxes[:, 1] = bboxes[:, 1] + bboxes[:, 3] * 0.5
return bboxes