| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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.bbox_utils import nonempty_bbox |
| from .transformers import bbox_cxcywh_to_xyxy |
| try: |
| from collections.abc import Sequence |
| except Exception: |
| from collections import Sequence |
|
|
| __all__ = [ |
| 'BBoxPostProcess', 'MaskPostProcess', 'JDEBBoxPostProcess', |
| 'CenterNetPostProcess', 'DETRBBoxPostProcess', 'SparsePostProcess' |
| ] |
|
|
|
|
| @register |
| class BBoxPostProcess(object): |
| __shared__ = ['num_classes', 'export_onnx', 'export_eb'] |
| __inject__ = ['decode', 'nms'] |
|
|
| def __init__(self, |
| num_classes=80, |
| decode=None, |
| nms=None, |
| export_onnx=False, |
| export_eb=False): |
| super(BBoxPostProcess, self).__init__() |
| self.num_classes = num_classes |
| self.decode = decode |
| self.nms = nms |
| self.export_onnx = export_onnx |
| self.export_eb = export_eb |
|
|
| def __call__(self, head_out, rois, im_shape, scale_factor): |
| """ |
| Decode the bbox and do NMS if needed. |
| |
| Args: |
| head_out (tuple): bbox_pred and cls_prob of bbox_head output. |
| rois (tuple): roi and rois_num of rpn_head output. |
| im_shape (Tensor): The shape of the input image. |
| scale_factor (Tensor): The scale factor of the input image. |
| export_onnx (bool): whether export model to onnx |
| Returns: |
| bbox_pred (Tensor): The output prediction with shape [N, 6], including |
| labels, scores and bboxes. The size of bboxes are corresponding |
| to the input image, the bboxes may be used in other branch. |
| bbox_num (Tensor): The number of prediction boxes of each batch with |
| shape [1], and is N. |
| """ |
| if self.nms is not None: |
| bboxes, score = self.decode(head_out, rois, im_shape, scale_factor) |
| bbox_pred, bbox_num, before_nms_indexes = self.nms(bboxes, score, self.num_classes) |
|
|
| else: |
| bbox_pred, bbox_num = self.decode(head_out, rois, im_shape, |
| scale_factor) |
|
|
| if self.export_onnx: |
| |
| fake_bboxes = paddle.to_tensor( |
| np.array( |
| [[0., 0.0, 0.0, 0.0, 1.0, 1.0]], dtype='float32')) |
|
|
| bbox_pred = paddle.concat([bbox_pred, fake_bboxes]) |
| bbox_num = bbox_num + 1 |
|
|
| if self.nms is not None: |
| return bbox_pred, bbox_num, before_nms_indexes |
| else: |
| return bbox_pred, bbox_num |
|
|
| def get_pred(self, bboxes, bbox_num, im_shape, scale_factor): |
| """ |
| Rescale, clip and filter the bbox from the output of NMS to |
| get final prediction. |
| |
| Notes: |
| Currently only support bs = 1. |
| |
| Args: |
| bboxes (Tensor): The output bboxes with shape [N, 6] after decode |
| and NMS, including labels, scores and bboxes. |
| bbox_num (Tensor): The number of prediction boxes of each batch with |
| shape [1], and is N. |
| im_shape (Tensor): The shape of the input image. |
| scale_factor (Tensor): The scale factor of the input image. |
| Returns: |
| pred_result (Tensor): The final prediction results with shape [N, 6] |
| including labels, scores and bboxes. |
| """ |
| if self.export_eb: |
| |
| return bboxes, bboxes, bbox_num |
|
|
| if not self.export_onnx: |
| bboxes_list = [] |
| bbox_num_list = [] |
| id_start = 0 |
| fake_bboxes = paddle.to_tensor( |
| np.array( |
| [[0., 0.0, 0.0, 0.0, 1.0, 1.0]], dtype='float32')) |
| fake_bbox_num = paddle.to_tensor(np.array([1], dtype='int32')) |
|
|
| |
| for i in range(bbox_num.shape[0]): |
| if bbox_num[i] == 0: |
| bboxes_i = fake_bboxes |
| bbox_num_i = fake_bbox_num |
| else: |
| bboxes_i = bboxes[id_start:id_start + bbox_num[i], :] |
| bbox_num_i = bbox_num[i] |
| id_start += bbox_num[i] |
| bboxes_list.append(bboxes_i) |
| bbox_num_list.append(bbox_num_i) |
| bboxes = paddle.concat(bboxes_list) |
| bbox_num = paddle.concat(bbox_num_list) |
|
|
| origin_shape = paddle.floor(im_shape / scale_factor + 0.5) |
|
|
| if not self.export_onnx: |
| origin_shape_list = [] |
| scale_factor_list = [] |
| |
| for i in range(bbox_num.shape[0]): |
| expand_shape = paddle.expand(origin_shape[i:i + 1, :], |
| [bbox_num[i], 2]) |
| scale_y, scale_x = scale_factor[i][0], scale_factor[i][1] |
| scale = paddle.concat([scale_x, scale_y, scale_x, scale_y]) |
| expand_scale = paddle.expand(scale, [bbox_num[i], 4]) |
| origin_shape_list.append(expand_shape) |
| scale_factor_list.append(expand_scale) |
|
|
| self.origin_shape_list = paddle.concat(origin_shape_list) |
| scale_factor_list = paddle.concat(scale_factor_list) |
|
|
| else: |
| |
| scale_y, scale_x = scale_factor[0][0], scale_factor[0][1] |
| scale = paddle.concat( |
| [scale_x, scale_y, scale_x, scale_y]).unsqueeze(0) |
| self.origin_shape_list = paddle.expand(origin_shape, |
| [bbox_num[0], 2]) |
| scale_factor_list = paddle.expand(scale, [bbox_num[0], 4]) |
|
|
| |
| pred_label = bboxes[:, 0:1] |
| pred_score = bboxes[:, 1:2] |
| pred_bbox = bboxes[:, 2:] |
| |
| scaled_bbox = pred_bbox / scale_factor_list |
| origin_h = self.origin_shape_list[:, 0] |
| origin_w = self.origin_shape_list[:, 1] |
| zeros = paddle.zeros_like(origin_h) |
| |
| x1 = paddle.maximum(paddle.minimum(scaled_bbox[:, 0], origin_w), zeros) |
| y1 = paddle.maximum(paddle.minimum(scaled_bbox[:, 1], origin_h), zeros) |
| x2 = paddle.maximum(paddle.minimum(scaled_bbox[:, 2], origin_w), zeros) |
| y2 = paddle.maximum(paddle.minimum(scaled_bbox[:, 3], origin_h), zeros) |
| pred_bbox = paddle.stack([x1, y1, x2, y2], axis=-1) |
| |
| keep_mask = nonempty_bbox(pred_bbox, return_mask=True) |
| keep_mask = paddle.unsqueeze(keep_mask, [1]) |
| pred_label = paddle.where(keep_mask, pred_label, |
| paddle.ones_like(pred_label) * -1) |
| pred_result = paddle.concat([pred_label, pred_score, pred_bbox], axis=1) |
| return bboxes, pred_result, bbox_num |
|
|
| def get_origin_shape(self, ): |
| return self.origin_shape_list |
|
|
|
|
| @register |
| class MaskPostProcess(object): |
| __shared__ = ['export_onnx', 'assign_on_cpu'] |
| """ |
| refer to: |
| https://github.com/facebookresearch/detectron2/layers/mask_ops.py |
| |
| Get Mask output according to the output from model |
| """ |
|
|
| def __init__(self, |
| binary_thresh=0.5, |
| export_onnx=False, |
| assign_on_cpu=False): |
| super(MaskPostProcess, self).__init__() |
| self.binary_thresh = binary_thresh |
| self.export_onnx = export_onnx |
| self.assign_on_cpu = assign_on_cpu |
|
|
| def __call__(self, mask_out, bboxes, bbox_num, origin_shape): |
| """ |
| Decode the mask_out and paste the mask to the origin image. |
| |
| Args: |
| mask_out (Tensor): mask_head output with shape [N, 28, 28]. |
| bbox_pred (Tensor): The output bboxes with shape [N, 6] after decode |
| and NMS, including labels, scores and bboxes. |
| bbox_num (Tensor): The number of prediction boxes of each batch with |
| shape [1], and is N. |
| origin_shape (Tensor): The origin shape of the input image, the tensor |
| shape is [N, 2], and each row is [h, w]. |
| Returns: |
| pred_result (Tensor): The final prediction mask results with shape |
| [N, h, w] in binary mask style. |
| """ |
| num_mask = mask_out.shape[0] |
| origin_shape = paddle.cast(origin_shape, 'int32') |
| device = paddle.device.get_device() |
|
|
| if self.export_onnx: |
| h, w = origin_shape[0][0], origin_shape[0][1] |
| mask_onnx = paste_mask(mask_out[:, None, :, :], bboxes[:, 2:], h, w, |
| self.assign_on_cpu) |
| mask_onnx = mask_onnx >= self.binary_thresh |
| pred_result = paddle.cast(mask_onnx, 'int32') |
|
|
| else: |
| max_h = paddle.max(origin_shape[:, 0]) |
| max_w = paddle.max(origin_shape[:, 1]) |
| pred_result = paddle.zeros( |
| [num_mask, max_h, max_w], dtype='int32') - 1 |
|
|
| id_start = 0 |
| for i in range(paddle.shape(bbox_num)[0]): |
| bboxes_i = bboxes[id_start:id_start + bbox_num[i], :] |
| mask_out_i = mask_out[id_start:id_start + bbox_num[i], :, :] |
| im_h = origin_shape[i, 0] |
| im_w = origin_shape[i, 1] |
| pred_mask = paste_mask(mask_out_i[:, None, :, :], |
| bboxes_i[:, 2:], im_h, im_w, |
| self.assign_on_cpu) |
| pred_mask = paddle.cast(pred_mask >= self.binary_thresh, |
| 'int32') |
| pred_result[id_start:id_start + bbox_num[i], :im_h, : |
| im_w] = pred_mask |
| id_start += bbox_num[i] |
| if self.assign_on_cpu: |
| paddle.set_device(device) |
|
|
| return pred_result |
|
|
|
|
| @register |
| class JDEBBoxPostProcess(nn.Layer): |
| __shared__ = ['num_classes'] |
| __inject__ = ['decode', 'nms'] |
|
|
| def __init__(self, num_classes=1, decode=None, nms=None, return_idx=True): |
| super(JDEBBoxPostProcess, self).__init__() |
| self.num_classes = num_classes |
| self.decode = decode |
| self.nms = nms |
| self.return_idx = return_idx |
|
|
| self.fake_bbox_pred = paddle.to_tensor( |
| np.array( |
| [[-1, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype='float32')) |
| self.fake_bbox_num = paddle.to_tensor(np.array([1], dtype='int32')) |
| self.fake_nms_keep_idx = paddle.to_tensor( |
| np.array( |
| [[0]], dtype='int32')) |
|
|
| self.fake_yolo_boxes_out = paddle.to_tensor( |
| np.array( |
| [[[0.0, 0.0, 0.0, 0.0]]], dtype='float32')) |
| self.fake_yolo_scores_out = paddle.to_tensor( |
| np.array( |
| [[[0.0]]], dtype='float32')) |
| self.fake_boxes_idx = paddle.to_tensor(np.array([[0]], dtype='int64')) |
|
|
| def forward(self, head_out, anchors): |
| """ |
| Decode the bbox and do NMS for JDE model. |
| |
| Args: |
| head_out (list): Bbox_pred and cls_prob of bbox_head output. |
| anchors (list): Anchors of JDE model. |
| |
| Returns: |
| boxes_idx (Tensor): The index of kept bboxes after decode 'JDEBox'. |
| bbox_pred (Tensor): The output is the prediction with shape [N, 6] |
| including labels, scores and bboxes. |
| bbox_num (Tensor): The number of prediction of each batch with shape [N]. |
| nms_keep_idx (Tensor): The index of kept bboxes after NMS. |
| """ |
| boxes_idx, yolo_boxes_scores = self.decode(head_out, anchors) |
|
|
| if len(boxes_idx) == 0: |
| boxes_idx = self.fake_boxes_idx |
| yolo_boxes_out = self.fake_yolo_boxes_out |
| yolo_scores_out = self.fake_yolo_scores_out |
| else: |
| yolo_boxes = paddle.gather_nd(yolo_boxes_scores, boxes_idx) |
| |
| yolo_boxes_out = paddle.reshape( |
| yolo_boxes[:, :4], shape=[1, len(boxes_idx), 4]) |
| yolo_scores_out = paddle.reshape( |
| yolo_boxes[:, 4:5], shape=[1, 1, len(boxes_idx)]) |
| boxes_idx = boxes_idx[:, 1:] |
|
|
| if self.return_idx: |
| bbox_pred, bbox_num, nms_keep_idx = self.nms( |
| yolo_boxes_out, yolo_scores_out, self.num_classes) |
| if bbox_pred.shape[0] == 0: |
| bbox_pred = self.fake_bbox_pred |
| bbox_num = self.fake_bbox_num |
| nms_keep_idx = self.fake_nms_keep_idx |
| return boxes_idx, bbox_pred, bbox_num, nms_keep_idx |
| else: |
| bbox_pred, bbox_num, _ = self.nms(yolo_boxes_out, yolo_scores_out, |
| self.num_classes) |
| if bbox_pred.shape[0] == 0: |
| bbox_pred = self.fake_bbox_pred |
| bbox_num = self.fake_bbox_num |
| return _, bbox_pred, bbox_num, _ |
|
|
|
|
| @register |
| class CenterNetPostProcess(object): |
| """ |
| Postprocess the model outputs to get final prediction: |
| 1. Do NMS for heatmap to get top `max_per_img` bboxes. |
| 2. Decode bboxes using center offset and box size. |
| 3. Rescale decoded bboxes reference to the origin image shape. |
| Args: |
| max_per_img(int): the maximum number of predicted objects in a image, |
| 500 by default. |
| down_ratio(int): the down ratio from images to heatmap, 4 by default. |
| regress_ltrb (bool): whether to regress left/top/right/bottom or |
| width/height for a box, true by default. |
| """ |
| __shared__ = ['down_ratio'] |
|
|
| def __init__(self, max_per_img=500, down_ratio=4, regress_ltrb=True): |
| super(CenterNetPostProcess, self).__init__() |
| self.max_per_img = max_per_img |
| self.down_ratio = down_ratio |
| self.regress_ltrb = regress_ltrb |
| |
|
|
| def _simple_nms(self, heat, kernel=3): |
| """ Use maxpool to filter the max score, get local peaks. """ |
| pad = (kernel - 1) // 2 |
| hmax = F.max_pool2d(heat, kernel, stride=1, padding=pad) |
| keep = paddle.cast(hmax == heat, 'float32') |
| return heat * keep |
|
|
| def _topk(self, scores): |
| """ Select top k scores and decode to get xy coordinates. """ |
| k = self.max_per_img |
| shape_fm = paddle.shape(scores) |
| shape_fm.stop_gradient = True |
| cat, height, width = shape_fm[1], shape_fm[2], shape_fm[3] |
| |
| scores_r = paddle.reshape(scores, [cat, -1]) |
| topk_scores, topk_inds = paddle.topk(scores_r, k) |
| topk_ys = topk_inds // width |
| topk_xs = topk_inds % width |
|
|
| topk_score_r = paddle.reshape(topk_scores, [-1]) |
| topk_score, topk_ind = paddle.topk(topk_score_r, k) |
| k_t = paddle.full(paddle.shape(topk_ind), k, dtype='int64') |
| topk_clses = paddle.cast(paddle.floor_divide(topk_ind, k_t), 'float32') |
|
|
| topk_inds = paddle.reshape(topk_inds, [-1]) |
| topk_ys = paddle.reshape(topk_ys, [-1, 1]) |
| topk_xs = paddle.reshape(topk_xs, [-1, 1]) |
| topk_inds = paddle.gather(topk_inds, topk_ind) |
| topk_ys = paddle.gather(topk_ys, topk_ind) |
| topk_xs = paddle.gather(topk_xs, topk_ind) |
| return topk_score, topk_inds, topk_clses, topk_ys, topk_xs |
|
|
| def __call__(self, hm, wh, reg, im_shape, scale_factor): |
| |
| heat = self._simple_nms(hm) |
| scores, inds, topk_clses, ys, xs = self._topk(heat) |
| clses = topk_clses.unsqueeze(1) |
| scores = scores.unsqueeze(1) |
|
|
| |
| reg_t = paddle.transpose(reg, [0, 2, 3, 1]) |
| reg = paddle.reshape(reg_t, [-1, reg_t.shape[-1]]) |
| reg = paddle.gather(reg, inds) |
| xs = paddle.cast(xs, 'float32') |
| ys = paddle.cast(ys, 'float32') |
| xs = xs + reg[:, 0:1] |
| ys = ys + reg[:, 1:2] |
| wh_t = paddle.transpose(wh, [0, 2, 3, 1]) |
| wh = paddle.reshape(wh_t, [-1, wh_t.shape[-1]]) |
| wh = paddle.gather(wh, inds) |
| if self.regress_ltrb: |
| x1 = xs - wh[:, 0:1] |
| y1 = ys - wh[:, 1:2] |
| x2 = xs + wh[:, 2:3] |
| y2 = ys + wh[:, 3:4] |
| else: |
| x1 = xs - wh[:, 0:1] / 2 |
| y1 = ys - wh[:, 1:2] / 2 |
| x2 = xs + wh[:, 0:1] / 2 |
| y2 = ys + wh[:, 1:2] / 2 |
| n, c, feat_h, feat_w = paddle.shape(hm) |
| padw = (feat_w * self.down_ratio - im_shape[0, 1]) / 2 |
| padh = (feat_h * self.down_ratio - im_shape[0, 0]) / 2 |
| x1 = x1 * self.down_ratio |
| y1 = y1 * self.down_ratio |
| x2 = x2 * self.down_ratio |
| y2 = y2 * self.down_ratio |
| x1 = x1 - padw |
| y1 = y1 - padh |
| x2 = x2 - padw |
| y2 = y2 - padh |
| bboxes = paddle.concat([x1, y1, x2, y2], axis=1) |
| scale_y = scale_factor[:, 0:1] |
| scale_x = scale_factor[:, 1:2] |
| scale_expand = paddle.concat( |
| [scale_x, scale_y, scale_x, scale_y], axis=1) |
| boxes_shape = bboxes.shape[:] |
| scale_expand = paddle.expand(scale_expand, shape=boxes_shape) |
| bboxes = paddle.divide(bboxes, scale_expand) |
|
|
| results = paddle.concat([clses, scores, bboxes], axis=1) |
| return results, paddle.shape(results)[0:1], inds, topk_clses, ys, xs |
|
|
|
|
| @register |
| class DETRBBoxPostProcess(object): |
| __shared__ = ['num_classes', 'use_focal_loss'] |
| __inject__ = [] |
|
|
| def __init__(self, |
| num_classes=80, |
| num_top_queries=100, |
| use_focal_loss=False): |
| super(DETRBBoxPostProcess, self).__init__() |
| self.num_classes = num_classes |
| self.num_top_queries = num_top_queries |
| self.use_focal_loss = use_focal_loss |
|
|
| def __call__(self, head_out, im_shape, scale_factor): |
| """ |
| Decode the bbox. |
| |
| Args: |
| head_out (tuple): bbox_pred, cls_logit and masks of bbox_head output. |
| im_shape (Tensor): The shape of the input image. |
| scale_factor (Tensor): The scale factor of the input image. |
| Returns: |
| bbox_pred (Tensor): The output prediction with shape [N, 6], including |
| labels, scores and bboxes. The size of bboxes are corresponding |
| to the input image, the bboxes may be used in other branch. |
| bbox_num (Tensor): The number of prediction boxes of each batch with |
| shape [bs], and is N. |
| """ |
| bboxes, logits, masks = head_out |
|
|
| bbox_pred = bbox_cxcywh_to_xyxy(bboxes) |
| origin_shape = paddle.floor(im_shape / scale_factor + 0.5) |
| img_h, img_w = paddle.split(origin_shape, 2, axis=-1) |
| origin_shape = paddle.concat( |
| [img_w, img_h, img_w, img_h], axis=-1).reshape([-1, 1, 4]) |
| bbox_pred *= origin_shape |
|
|
| scores = F.sigmoid(logits) if self.use_focal_loss else F.softmax( |
| logits)[:, :, :-1] |
|
|
| if not self.use_focal_loss: |
| scores, labels = scores.max(-1), scores.argmax(-1) |
| if scores.shape[1] > self.num_top_queries: |
| scores, index = paddle.topk( |
| scores, self.num_top_queries, axis=-1) |
| batch_ind = paddle.arange( |
| end=scores.shape[0]).unsqueeze(-1).tile( |
| [1, self.num_top_queries]) |
| index = paddle.stack([batch_ind, index], axis=-1) |
| labels = paddle.gather_nd(labels, index) |
| bbox_pred = paddle.gather_nd(bbox_pred, index) |
| else: |
| scores, index = paddle.topk( |
| scores.flatten(1), self.num_top_queries, axis=-1) |
| labels = index % self.num_classes |
| index = index // self.num_classes |
| batch_ind = paddle.arange(end=scores.shape[0]).unsqueeze(-1).tile( |
| [1, self.num_top_queries]) |
| index = paddle.stack([batch_ind, index], axis=-1) |
| bbox_pred = paddle.gather_nd(bbox_pred, index) |
|
|
| bbox_pred = paddle.concat( |
| [ |
| labels.unsqueeze(-1).astype('float32'), scores.unsqueeze(-1), |
| bbox_pred |
| ], |
| axis=-1) |
| bbox_num = paddle.to_tensor( |
| bbox_pred.shape[1], dtype='int32').tile([bbox_pred.shape[0]]) |
| bbox_pred = bbox_pred.reshape([-1, 6]) |
| return bbox_pred, bbox_num |
|
|
|
|
| @register |
| class SparsePostProcess(object): |
| __shared__ = ['num_classes', 'assign_on_cpu'] |
|
|
| def __init__(self, |
| num_proposals, |
| num_classes=80, |
| binary_thresh=0.5, |
| assign_on_cpu=False): |
| super(SparsePostProcess, self).__init__() |
| self.num_classes = num_classes |
| self.num_proposals = num_proposals |
| self.binary_thresh = binary_thresh |
| self.assign_on_cpu = assign_on_cpu |
|
|
| def __call__(self, scores, bboxes, scale_factor, ori_shape, masks=None): |
| assert len(scores) == len(bboxes) == \ |
| len(ori_shape) == len(scale_factor) |
| device = paddle.device.get_device() |
| batch_size = len(ori_shape) |
|
|
| scores = F.sigmoid(scores) |
| has_mask = masks is not None |
| if has_mask: |
| masks = F.sigmoid(masks) |
| masks = masks.reshape([batch_size, -1, *masks.shape[1:]]) |
|
|
| bbox_pred = [] |
| mask_pred = [] if has_mask else None |
| bbox_num = paddle.zeros([batch_size], dtype='int32') |
| for i in range(batch_size): |
| score = scores[i] |
| bbox = bboxes[i] |
| score, indices = score.flatten(0, 1).topk( |
| self.num_proposals, sorted=False) |
| label = indices % self.num_classes |
| if has_mask: |
| mask = masks[i] |
| mask = mask.flatten(0, 1)[indices] |
|
|
| H, W = ori_shape[i][0], ori_shape[i][1] |
| bbox = bbox[paddle.cast(indices / self.num_classes, indices.dtype)] |
| bbox /= scale_factor[i] |
| bbox[:, 0::2] = paddle.clip(bbox[:, 0::2], 0, W) |
| bbox[:, 1::2] = paddle.clip(bbox[:, 1::2], 0, H) |
|
|
| keep = ((bbox[:, 2] - bbox[:, 0]).numpy() > 1.) & \ |
| ((bbox[:, 3] - bbox[:, 1]).numpy() > 1.) |
| if keep.sum() == 0: |
| bbox = paddle.zeros([1, 6], dtype='float32') |
| if has_mask: |
| mask = paddle.zeros([1, H, W], dtype='uint8') |
| else: |
| label = paddle.to_tensor(label.numpy()[keep]).astype( |
| 'float32').unsqueeze(-1) |
| score = paddle.to_tensor(score.numpy()[keep]).astype( |
| 'float32').unsqueeze(-1) |
| bbox = paddle.to_tensor(bbox.numpy()[keep]).astype('float32') |
| if has_mask: |
| mask = paddle.to_tensor(mask.numpy()[keep]).astype( |
| 'float32').unsqueeze(1) |
| mask = paste_mask(mask, bbox, H, W, self.assign_on_cpu) |
| mask = paddle.cast(mask >= self.binary_thresh, 'uint8') |
| bbox = paddle.concat([label, score, bbox], axis=-1) |
|
|
| bbox_num[i] = bbox.shape[0] |
| bbox_pred.append(bbox) |
| if has_mask: |
| mask_pred.append(mask) |
|
|
| bbox_pred = paddle.concat(bbox_pred) |
| mask_pred = paddle.concat(mask_pred) if has_mask else None |
|
|
| if self.assign_on_cpu: |
| paddle.set_device(device) |
|
|
| if has_mask: |
| return bbox_pred, bbox_num, mask_pred |
| else: |
| return bbox_pred, bbox_num |
|
|
|
|
| def paste_mask(masks, boxes, im_h, im_w, assign_on_cpu=False): |
| """ |
| Paste the mask prediction to the original image. |
| """ |
| x0_int, y0_int = 0, 0 |
| x1_int, y1_int = im_w, im_h |
| x0, y0, x1, y1 = paddle.split(boxes, 4, axis=1) |
| N = masks.shape[0] |
| img_y = paddle.arange(y0_int, y1_int) + 0.5 |
| img_x = paddle.arange(x0_int, x1_int) + 0.5 |
|
|
| img_y = (img_y - y0) / (y1 - y0) * 2 - 1 |
| img_x = (img_x - x0) / (x1 - x0) * 2 - 1 |
| |
|
|
| if assign_on_cpu: |
| paddle.set_device('cpu') |
| gx = img_x[:, None, :].expand( |
| [N, paddle.shape(img_y)[1], paddle.shape(img_x)[1]]) |
| gy = img_y[:, :, None].expand( |
| [N, paddle.shape(img_y)[1], paddle.shape(img_x)[1]]) |
| grid = paddle.stack([gx, gy], axis=3) |
| img_masks = F.grid_sample(masks, grid, align_corners=False) |
| return img_masks[:, 0] |
|
|
|
|
| def multiclass_nms(bboxs, num_classes, match_threshold=0.6, match_metric='iou'): |
| final_boxes = [] |
| for c in range(num_classes): |
| idxs = bboxs[:, 0] == c |
| if np.count_nonzero(idxs) == 0: continue |
| r = nms(bboxs[idxs, 1:], match_threshold, match_metric) |
| final_boxes.append(np.concatenate([np.full((r.shape[0], 1), c), r], 1)) |
| return final_boxes |
|
|
|
|
| def nms(dets, match_threshold=0.6, match_metric='iou'): |
| """ Apply NMS to avoid detecting too many overlapping bounding boxes. |
| Args: |
| dets: shape [N, 5], [score, x1, y1, x2, y2] |
| match_metric: 'iou' or 'ios' |
| match_threshold: overlap thresh for match metric. |
| """ |
| if dets.shape[0] == 0: |
| return dets[[], :] |
| scores = dets[:, 0] |
| x1 = dets[:, 1] |
| y1 = dets[:, 2] |
| x2 = dets[:, 3] |
| y2 = dets[:, 4] |
| areas = (x2 - x1 + 1) * (y2 - y1 + 1) |
| order = scores.argsort()[::-1] |
|
|
| ndets = dets.shape[0] |
| suppressed = np.zeros((ndets), dtype=np.int32) |
|
|
| for _i in range(ndets): |
| i = order[_i] |
| if suppressed[i] == 1: |
| continue |
| ix1 = x1[i] |
| iy1 = y1[i] |
| ix2 = x2[i] |
| iy2 = y2[i] |
| iarea = areas[i] |
| for _j in range(_i + 1, ndets): |
| j = order[_j] |
| if suppressed[j] == 1: |
| continue |
| xx1 = max(ix1, x1[j]) |
| yy1 = max(iy1, y1[j]) |
| xx2 = min(ix2, x2[j]) |
| yy2 = min(iy2, y2[j]) |
| w = max(0.0, xx2 - xx1 + 1) |
| h = max(0.0, yy2 - yy1 + 1) |
| inter = w * h |
| if match_metric == 'iou': |
| union = iarea + areas[j] - inter |
| match_value = inter / union |
| elif match_metric == 'ios': |
| smaller = min(iarea, areas[j]) |
| match_value = inter / smaller |
| else: |
| raise ValueError() |
| if match_value >= match_threshold: |
| suppressed[j] = 1 |
| keep = np.where(suppressed == 0)[0] |
| dets = dets[keep, :] |
| return dets |
|
|
