post_process.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import numpy as np
  15. import paddle
  16. import paddle.nn as nn
  17. import paddle.nn.functional as F
  18. from paddlex.ppdet.core.workspace import register
  19. from paddlex.ppdet.modeling.bbox_utils import nonempty_bbox, rbox2poly
  20. from paddlex.ppdet.modeling.layers import TTFBox
  21. from .transformers import bbox_cxcywh_to_xyxy
  22. try:
  23. from collections.abc import Sequence
  24. except Exception:
  25. from collections import Sequence
  26. __all__ = [
  27. 'BBoxPostProcess', 'MaskPostProcess', 'FCOSPostProcess',
  28. 'S2ANetBBoxPostProcess', 'JDEBBoxPostProcess', 'CenterNetPostProcess',
  29. 'DETRBBoxPostProcess', 'SparsePostProcess'
  30. ]
  31. @register
  32. class BBoxPostProcess(nn.Layer):
  33. __shared__ = ['num_classes']
  34. __inject__ = ['decode', 'nms']
  35. def __init__(self, num_classes=80, decode=None, nms=None):
  36. super(BBoxPostProcess, self).__init__()
  37. self.num_classes = num_classes
  38. self.decode = decode
  39. self.nms = nms
  40. self.fake_bboxes = paddle.to_tensor(
  41. np.array(
  42. [[-1, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype='float32'))
  43. self.fake_bbox_num = paddle.to_tensor(np.array([1], dtype='int32'))
  44. def forward(self, head_out, rois, im_shape, scale_factor):
  45. """
  46. Decode the bbox and do NMS if needed.
  47. Args:
  48. head_out (tuple): bbox_pred and cls_prob of bbox_head output.
  49. rois (tuple): roi and rois_num of rpn_head output.
  50. im_shape (Tensor): The shape of the input image.
  51. scale_factor (Tensor): The scale factor of the input image.
  52. Returns:
  53. bbox_pred (Tensor): The output prediction with shape [N, 6], including
  54. labels, scores and bboxes. The size of bboxes are corresponding
  55. to the input image, the bboxes may be used in other branch.
  56. bbox_num (Tensor): The number of prediction boxes of each batch with
  57. shape [1], and is N.
  58. """
  59. if self.nms is not None:
  60. bboxes, score = self.decode(head_out, rois, im_shape, scale_factor)
  61. bbox_pred, bbox_num, _ = self.nms(bboxes, score, self.num_classes)
  62. else:
  63. bbox_pred, bbox_num = self.decode(head_out, rois, im_shape,
  64. scale_factor)
  65. return bbox_pred, bbox_num
  66. def get_pred(self, bboxes, bbox_num, im_shape, scale_factor):
  67. """
  68. Rescale, clip and filter the bbox from the output of NMS to
  69. get final prediction.
  70. Notes:
  71. Currently only support bs = 1.
  72. Args:
  73. bboxes (Tensor): The output bboxes with shape [N, 6] after decode
  74. and NMS, including labels, scores and bboxes.
  75. bbox_num (Tensor): The number of prediction boxes of each batch with
  76. shape [1], and is N.
  77. im_shape (Tensor): The shape of the input image.
  78. scale_factor (Tensor): The scale factor of the input image.
  79. Returns:
  80. pred_result (Tensor): The final prediction results with shape [N, 6]
  81. including labels, scores and bboxes.
  82. """
  83. if bboxes.shape[0] == 0:
  84. bboxes = self.fake_bboxes
  85. bbox_num = self.fake_bbox_num
  86. origin_shape = paddle.floor(im_shape / scale_factor + 0.5)
  87. origin_shape_list = []
  88. scale_factor_list = []
  89. # scale_factor: scale_y, scale_x
  90. for i in range(bbox_num.shape[0]):
  91. expand_shape = paddle.expand(origin_shape[i:i + 1, :],
  92. [bbox_num[i], 2])
  93. scale_y, scale_x = scale_factor[i][0], scale_factor[i][1]
  94. scale = paddle.concat([scale_x, scale_y, scale_x, scale_y])
  95. expand_scale = paddle.expand(scale, [bbox_num[i], 4])
  96. origin_shape_list.append(expand_shape)
  97. scale_factor_list.append(expand_scale)
  98. self.origin_shape_list = paddle.concat(origin_shape_list)
  99. scale_factor_list = paddle.concat(scale_factor_list)
  100. # bboxes: [N, 6], label, score, bbox
  101. pred_label = bboxes[:, 0:1]
  102. pred_score = bboxes[:, 1:2]
  103. pred_bbox = bboxes[:, 2:]
  104. # rescale bbox to original image
  105. scaled_bbox = pred_bbox / scale_factor_list
  106. origin_h = self.origin_shape_list[:, 0]
  107. origin_w = self.origin_shape_list[:, 1]
  108. zeros = paddle.zeros_like(origin_h)
  109. # clip bbox to [0, original_size]
  110. x1 = paddle.maximum(paddle.minimum(scaled_bbox[:, 0], origin_w), zeros)
  111. y1 = paddle.maximum(paddle.minimum(scaled_bbox[:, 1], origin_h), zeros)
  112. x2 = paddle.maximum(paddle.minimum(scaled_bbox[:, 2], origin_w), zeros)
  113. y2 = paddle.maximum(paddle.minimum(scaled_bbox[:, 3], origin_h), zeros)
  114. pred_bbox = paddle.stack([x1, y1, x2, y2], axis=-1)
  115. # filter empty bbox
  116. keep_mask = nonempty_bbox(pred_bbox, return_mask=True)
  117. keep_mask = paddle.unsqueeze(keep_mask, [1])
  118. pred_label = paddle.where(keep_mask, pred_label,
  119. paddle.ones_like(pred_label) * -1)
  120. pred_result = paddle.concat(
  121. [pred_label, pred_score, pred_bbox], axis=1)
  122. return pred_result
  123. def get_origin_shape(self, ):
  124. return self.origin_shape_list
  125. @register
  126. class MaskPostProcess(object):
  127. def __init__(self, binary_thresh=0.5):
  128. super(MaskPostProcess, self).__init__()
  129. self.binary_thresh = binary_thresh
  130. def paste_mask(self, masks, boxes, im_h, im_w):
  131. """
  132. Paste the mask prediction to the original image.
  133. """
  134. x0, y0, x1, y1 = paddle.split(boxes, 4, axis=1)
  135. masks = paddle.unsqueeze(masks, [0, 1])
  136. img_y = paddle.arange(0, im_h, dtype='float32') + 0.5
  137. img_x = paddle.arange(0, im_w, dtype='float32') + 0.5
  138. img_y = (img_y - y0) / (y1 - y0) * 2 - 1
  139. img_x = (img_x - x0) / (x1 - x0) * 2 - 1
  140. img_x = paddle.unsqueeze(img_x, [1])
  141. img_y = paddle.unsqueeze(img_y, [2])
  142. N = boxes.shape[0]
  143. gx = paddle.expand(img_x, [N, img_y.shape[1], img_x.shape[2]])
  144. gy = paddle.expand(img_y, [N, img_y.shape[1], img_x.shape[2]])
  145. grid = paddle.stack([gx, gy], axis=3)
  146. img_masks = F.grid_sample(masks, grid, align_corners=False)
  147. return img_masks[:, 0]
  148. def __call__(self, mask_out, bboxes, bbox_num, origin_shape):
  149. """
  150. Decode the mask_out and paste the mask to the origin image.
  151. Args:
  152. mask_out (Tensor): mask_head output with shape [N, 28, 28].
  153. bbox_pred (Tensor): The output bboxes with shape [N, 6] after decode
  154. and NMS, including labels, scores and bboxes.
  155. bbox_num (Tensor): The number of prediction boxes of each batch with
  156. shape [1], and is N.
  157. origin_shape (Tensor): The origin shape of the input image, the tensor
  158. shape is [N, 2], and each row is [h, w].
  159. Returns:
  160. pred_result (Tensor): The final prediction mask results with shape
  161. [N, h, w] in binary mask style.
  162. """
  163. num_mask = mask_out.shape[0]
  164. origin_shape = paddle.cast(origin_shape, 'int32')
  165. # TODO: support bs > 1 and mask output dtype is bool
  166. pred_result = paddle.zeros(
  167. [num_mask, origin_shape[0][0], origin_shape[0][1]], dtype='int32')
  168. if bbox_num == 1 and bboxes[0][0] == -1:
  169. return pred_result
  170. # TODO: optimize chunk paste
  171. pred_result = []
  172. for i in range(bboxes.shape[0]):
  173. im_h, im_w = origin_shape[i][0], origin_shape[i][1]
  174. pred_mask = self.paste_mask(mask_out[i], bboxes[i:i + 1, 2:], im_h,
  175. im_w)
  176. pred_mask = pred_mask >= self.binary_thresh
  177. pred_mask = paddle.cast(pred_mask, 'int32')
  178. pred_result.append(pred_mask)
  179. pred_result = paddle.concat(pred_result)
  180. return pred_result
  181. @register
  182. class FCOSPostProcess(object):
  183. __inject__ = ['decode', 'nms']
  184. def __init__(self, decode=None, nms=None):
  185. super(FCOSPostProcess, self).__init__()
  186. self.decode = decode
  187. self.nms = nms
  188. def __call__(self, fcos_head_outs, scale_factor):
  189. """
  190. Decode the bbox and do NMS in FCOS.
  191. """
  192. locations, cls_logits, bboxes_reg, centerness = fcos_head_outs
  193. bboxes, score = self.decode(locations, cls_logits, bboxes_reg,
  194. centerness, scale_factor)
  195. bbox_pred, bbox_num, _ = self.nms(bboxes, score)
  196. return bbox_pred, bbox_num
  197. @register
  198. class S2ANetBBoxPostProcess(nn.Layer):
  199. __shared__ = ['num_classes']
  200. __inject__ = ['nms']
  201. def __init__(self, num_classes=15, nms_pre=2000, min_bbox_size=0,
  202. nms=None):
  203. super(S2ANetBBoxPostProcess, self).__init__()
  204. self.num_classes = num_classes
  205. self.nms_pre = paddle.to_tensor(nms_pre)
  206. self.min_bbox_size = min_bbox_size
  207. self.nms = nms
  208. self.origin_shape_list = []
  209. self.fake_pred_cls_score_bbox = paddle.to_tensor(
  210. np.array(
  211. [[-1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]],
  212. dtype='float32'))
  213. self.fake_bbox_num = paddle.to_tensor(np.array([1], dtype='int32'))
  214. def forward(self, pred_scores, pred_bboxes):
  215. """
  216. pred_scores : [N, M] score
  217. pred_bboxes : [N, 5] xc, yc, w, h, a
  218. im_shape : [N, 2] im_shape
  219. scale_factor : [N, 2] scale_factor
  220. """
  221. pred_ploys0 = rbox2poly(pred_bboxes)
  222. pred_ploys = paddle.unsqueeze(pred_ploys0, axis=0)
  223. # pred_scores [NA, 16] --> [16, NA]
  224. pred_scores0 = paddle.transpose(pred_scores, [1, 0])
  225. pred_scores = paddle.unsqueeze(pred_scores0, axis=0)
  226. pred_cls_score_bbox, bbox_num, _ = self.nms(pred_ploys, pred_scores,
  227. self.num_classes)
  228. # Prevent empty bbox_pred from decode or NMS.
  229. # Bboxes and score before NMS may be empty due to the score threshold.
  230. if pred_cls_score_bbox.shape[0] <= 0 or pred_cls_score_bbox.shape[
  231. 1] <= 1:
  232. pred_cls_score_bbox = self.fake_pred_cls_score_bbox
  233. bbox_num = self.fake_bbox_num
  234. pred_cls_score_bbox = paddle.reshape(pred_cls_score_bbox, [-1, 10])
  235. return pred_cls_score_bbox, bbox_num
  236. def get_pred(self, bboxes, bbox_num, im_shape, scale_factor):
  237. """
  238. Rescale, clip and filter the bbox from the output of NMS to
  239. get final prediction.
  240. Args:
  241. bboxes(Tensor): bboxes [N, 10]
  242. bbox_num(Tensor): bbox_num
  243. im_shape(Tensor): [1 2]
  244. scale_factor(Tensor): [1 2]
  245. Returns:
  246. bbox_pred(Tensor): The output is the prediction with shape [N, 8]
  247. including labels, scores and bboxes. The size of
  248. bboxes are corresponding to the original image.
  249. """
  250. origin_shape = paddle.floor(im_shape / scale_factor + 0.5)
  251. origin_shape_list = []
  252. scale_factor_list = []
  253. # scale_factor: scale_y, scale_x
  254. for i in range(bbox_num.shape[0]):
  255. expand_shape = paddle.expand(origin_shape[i:i + 1, :],
  256. [bbox_num[i], 2])
  257. scale_y, scale_x = scale_factor[i][0], scale_factor[i][1]
  258. scale = paddle.concat([
  259. scale_x, scale_y, scale_x, scale_y, scale_x, scale_y, scale_x,
  260. scale_y
  261. ])
  262. expand_scale = paddle.expand(scale, [bbox_num[i], 8])
  263. origin_shape_list.append(expand_shape)
  264. scale_factor_list.append(expand_scale)
  265. origin_shape_list = paddle.concat(origin_shape_list)
  266. scale_factor_list = paddle.concat(scale_factor_list)
  267. # bboxes: [N, 10], label, score, bbox
  268. pred_label_score = bboxes[:, 0:2]
  269. pred_bbox = bboxes[:, 2:]
  270. # rescale bbox to original image
  271. pred_bbox = pred_bbox.reshape([-1, 8])
  272. scaled_bbox = pred_bbox / scale_factor_list
  273. origin_h = origin_shape_list[:, 0]
  274. origin_w = origin_shape_list[:, 1]
  275. bboxes = scaled_bbox
  276. zeros = paddle.zeros_like(origin_h)
  277. x1 = paddle.maximum(paddle.minimum(bboxes[:, 0], origin_w - 1), zeros)
  278. y1 = paddle.maximum(paddle.minimum(bboxes[:, 1], origin_h - 1), zeros)
  279. x2 = paddle.maximum(paddle.minimum(bboxes[:, 2], origin_w - 1), zeros)
  280. y2 = paddle.maximum(paddle.minimum(bboxes[:, 3], origin_h - 1), zeros)
  281. x3 = paddle.maximum(paddle.minimum(bboxes[:, 4], origin_w - 1), zeros)
  282. y3 = paddle.maximum(paddle.minimum(bboxes[:, 5], origin_h - 1), zeros)
  283. x4 = paddle.maximum(paddle.minimum(bboxes[:, 6], origin_w - 1), zeros)
  284. y4 = paddle.maximum(paddle.minimum(bboxes[:, 7], origin_h - 1), zeros)
  285. pred_bbox = paddle.stack([x1, y1, x2, y2, x3, y3, x4, y4], axis=-1)
  286. pred_result = paddle.concat([pred_label_score, pred_bbox], axis=1)
  287. return pred_result
  288. @register
  289. class JDEBBoxPostProcess(nn.Layer):
  290. __shared__ = ['num_classes']
  291. __inject__ = ['decode', 'nms']
  292. def __init__(self, num_classes=1, decode=None, nms=None, return_idx=True):
  293. super(JDEBBoxPostProcess, self).__init__()
  294. self.num_classes = num_classes
  295. self.decode = decode
  296. self.nms = nms
  297. self.return_idx = return_idx
  298. self.fake_bbox_pred = paddle.to_tensor(
  299. np.array(
  300. [[-1, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype='float32'))
  301. self.fake_bbox_num = paddle.to_tensor(np.array([1], dtype='int32'))
  302. self.fake_nms_keep_idx = paddle.to_tensor(
  303. np.array(
  304. [[0]], dtype='int32'))
  305. self.fake_yolo_boxes_out = paddle.to_tensor(
  306. np.array(
  307. [[[0.0, 0.0, 0.0, 0.0]]], dtype='float32'))
  308. self.fake_yolo_scores_out = paddle.to_tensor(
  309. np.array(
  310. [[[0.0]]], dtype='float32'))
  311. self.fake_boxes_idx = paddle.to_tensor(np.array([[0]], dtype='int64'))
  312. def forward(self, head_out, anchors):
  313. """
  314. Decode the bbox and do NMS for JDE model.
  315. Args:
  316. head_out (list): Bbox_pred and cls_prob of bbox_head output.
  317. anchors (list): Anchors of JDE model.
  318. Returns:
  319. boxes_idx (Tensor): The index of kept bboxes after decode 'JDEBox'.
  320. bbox_pred (Tensor): The output is the prediction with shape [N, 6]
  321. including labels, scores and bboxes.
  322. bbox_num (Tensor): The number of prediction of each batch with shape [N].
  323. nms_keep_idx (Tensor): The index of kept bboxes after NMS.
  324. """
  325. boxes_idx, yolo_boxes_scores = self.decode(head_out, anchors)
  326. if len(boxes_idx) == 0:
  327. boxes_idx = self.fake_boxes_idx
  328. yolo_boxes_out = self.fake_yolo_boxes_out
  329. yolo_scores_out = self.fake_yolo_scores_out
  330. else:
  331. yolo_boxes = paddle.gather_nd(yolo_boxes_scores, boxes_idx)
  332. # TODO: only support bs=1 now
  333. yolo_boxes_out = paddle.reshape(
  334. yolo_boxes[:, :4], shape=[1, len(boxes_idx), 4])
  335. yolo_scores_out = paddle.reshape(
  336. yolo_boxes[:, 4:5], shape=[1, 1, len(boxes_idx)])
  337. boxes_idx = boxes_idx[:, 1:]
  338. if self.return_idx:
  339. bbox_pred, bbox_num, nms_keep_idx = self.nms(
  340. yolo_boxes_out, yolo_scores_out, self.num_classes)
  341. if bbox_pred.shape[0] == 0:
  342. bbox_pred = self.fake_bbox_pred
  343. bbox_num = self.fake_bbox_num
  344. nms_keep_idx = self.fake_nms_keep_idx
  345. return boxes_idx, bbox_pred, bbox_num, nms_keep_idx
  346. else:
  347. bbox_pred, bbox_num, _ = self.nms(yolo_boxes_out, yolo_scores_out,
  348. self.num_classes)
  349. if bbox_pred.shape[0] == 0:
  350. bbox_pred = self.fake_bbox_pred
  351. bbox_num = self.fake_bbox_num
  352. return _, bbox_pred, bbox_num, _
  353. @register
  354. class CenterNetPostProcess(TTFBox):
  355. """
  356. Postprocess the model outputs to get final prediction:
  357. 1. Do NMS for heatmap to get top `max_per_img` bboxes.
  358. 2. Decode bboxes using center offset and box size.
  359. 3. Rescale decoded bboxes reference to the origin image shape.
  360. Args:
  361. max_per_img(int): the maximum number of predicted objects in a image,
  362. 500 by default.
  363. down_ratio(int): the down ratio from images to heatmap, 4 by default.
  364. regress_ltrb (bool): whether to regress left/top/right/bottom or
  365. width/height for a box, true by default.
  366. for_mot (bool): whether return other features used in tracking model.
  367. """
  368. __shared__ = ['down_ratio']
  369. def __init__(self,
  370. max_per_img=500,
  371. down_ratio=4,
  372. regress_ltrb=True,
  373. for_mot=False):
  374. super(TTFBox, self).__init__()
  375. self.max_per_img = max_per_img
  376. self.down_ratio = down_ratio
  377. self.regress_ltrb = regress_ltrb
  378. self.for_mot = for_mot
  379. def __call__(self, hm, wh, reg, im_shape, scale_factor):
  380. heat = self._simple_nms(hm)
  381. scores, inds, clses, ys, xs = self._topk(heat)
  382. scores = paddle.tensor.unsqueeze(scores, [1])
  383. clses = paddle.tensor.unsqueeze(clses, [1])
  384. reg_t = paddle.transpose(reg, [0, 2, 3, 1])
  385. # Like TTFBox, batch size is 1.
  386. # TODO: support batch size > 1
  387. reg = paddle.reshape(reg_t, [-1, paddle.shape(reg_t)[-1]])
  388. reg = paddle.gather(reg, inds)
  389. xs = paddle.cast(xs, 'float32')
  390. ys = paddle.cast(ys, 'float32')
  391. xs = xs + reg[:, 0:1]
  392. ys = ys + reg[:, 1:2]
  393. wh_t = paddle.transpose(wh, [0, 2, 3, 1])
  394. wh = paddle.reshape(wh_t, [-1, paddle.shape(wh_t)[-1]])
  395. wh = paddle.gather(wh, inds)
  396. if self.regress_ltrb:
  397. x1 = xs - wh[:, 0:1]
  398. y1 = ys - wh[:, 1:2]
  399. x2 = xs + wh[:, 2:3]
  400. y2 = ys + wh[:, 3:4]
  401. else:
  402. x1 = xs - wh[:, 0:1] / 2
  403. y1 = ys - wh[:, 1:2] / 2
  404. x2 = xs + wh[:, 0:1] / 2
  405. y2 = ys + wh[:, 1:2] / 2
  406. n, c, feat_h, feat_w = hm.shape[:]
  407. padw = (feat_w * self.down_ratio - im_shape[0, 1]) / 2
  408. padh = (feat_h * self.down_ratio - im_shape[0, 0]) / 2
  409. x1 = x1 * self.down_ratio
  410. y1 = y1 * self.down_ratio
  411. x2 = x2 * self.down_ratio
  412. y2 = y2 * self.down_ratio
  413. x1 = x1 - padw
  414. y1 = y1 - padh
  415. x2 = x2 - padw
  416. y2 = y2 - padh
  417. bboxes = paddle.concat([x1, y1, x2, y2], axis=1)
  418. scale_y = scale_factor[:, 0:1]
  419. scale_x = scale_factor[:, 1:2]
  420. scale_expand = paddle.concat(
  421. [scale_x, scale_y, scale_x, scale_y], axis=1)
  422. boxes_shape = paddle.shape(bboxes)
  423. boxes_shape.stop_gradient = True
  424. scale_expand = paddle.expand(scale_expand, shape=boxes_shape)
  425. bboxes = paddle.divide(bboxes, scale_expand)
  426. if self.for_mot:
  427. results = paddle.concat([bboxes, scores, clses], axis=1)
  428. return results, inds
  429. else:
  430. results = paddle.concat([clses, scores, bboxes], axis=1)
  431. return results, paddle.shape(results)[0:1]
  432. @register
  433. class DETRBBoxPostProcess(object):
  434. __shared__ = ['num_classes', 'use_focal_loss']
  435. __inject__ = []
  436. def __init__(self,
  437. num_classes=80,
  438. num_top_queries=100,
  439. use_focal_loss=False):
  440. super(DETRBBoxPostProcess, self).__init__()
  441. self.num_classes = num_classes
  442. self.num_top_queries = num_top_queries
  443. self.use_focal_loss = use_focal_loss
  444. def __call__(self, head_out, im_shape, scale_factor):
  445. """
  446. Decode the bbox.
  447. Args:
  448. head_out (tuple): bbox_pred, cls_logit and masks of bbox_head output.
  449. im_shape (Tensor): The shape of the input image.
  450. scale_factor (Tensor): The scale factor of the input image.
  451. Returns:
  452. bbox_pred (Tensor): The output prediction with shape [N, 6], including
  453. labels, scores and bboxes. The size of bboxes are corresponding
  454. to the input image, the bboxes may be used in other branch.
  455. bbox_num (Tensor): The number of prediction boxes of each batch with
  456. shape [bs], and is N.
  457. """
  458. bboxes, logits, masks = head_out
  459. bbox_pred = bbox_cxcywh_to_xyxy(bboxes)
  460. origin_shape = paddle.floor(im_shape / scale_factor + 0.5)
  461. img_h, img_w = origin_shape.unbind(1)
  462. origin_shape = paddle.stack(
  463. [img_w, img_h, img_w, img_h], axis=-1).unsqueeze(0)
  464. bbox_pred *= origin_shape
  465. scores = F.sigmoid(logits) if self.use_focal_loss else F.softmax(
  466. logits)[:, :, :-1]
  467. if not self.use_focal_loss:
  468. scores, labels = scores.max(-1), scores.argmax(-1)
  469. if scores.shape[1] > self.num_top_queries:
  470. scores, index = paddle.topk(
  471. scores, self.num_top_queries, axis=-1)
  472. labels = paddle.stack(
  473. [paddle.gather(l, i) for l, i in zip(labels, index)])
  474. bbox_pred = paddle.stack(
  475. [paddle.gather(b, i) for b, i in zip(bbox_pred, index)])
  476. else:
  477. scores, index = paddle.topk(
  478. scores.reshape([logits.shape[0], -1]),
  479. self.num_top_queries,
  480. axis=-1)
  481. labels = index % logits.shape[2]
  482. index = index // logits.shape[2]
  483. bbox_pred = paddle.stack(
  484. [paddle.gather(b, i) for b, i in zip(bbox_pred, index)])
  485. bbox_pred = paddle.concat(
  486. [
  487. labels.unsqueeze(-1).astype('float32'), scores.unsqueeze(-1),
  488. bbox_pred
  489. ],
  490. axis=-1)
  491. bbox_num = paddle.to_tensor(
  492. bbox_pred.shape[1], dtype='int32').tile([bbox_pred.shape[0]])
  493. bbox_pred = bbox_pred.reshape([-1, 6])
  494. return bbox_pred, bbox_num
  495. @register
  496. class SparsePostProcess(object):
  497. __shared__ = ['num_classes']
  498. def __init__(self, num_proposals, num_classes=80):
  499. super(SparsePostProcess, self).__init__()
  500. self.num_classes = num_classes
  501. self.num_proposals = num_proposals
  502. def __call__(self, box_cls, box_pred, scale_factor_wh, img_whwh):
  503. """
  504. Arguments:
  505. box_cls (Tensor): tensor of shape (batch_size, num_proposals, K).
  506. The tensor predicts the classification probability for each proposal.
  507. box_pred (Tensor): tensors of shape (batch_size, num_proposals, 4).
  508. The tensor predicts 4-vector (x,y,w,h) box
  509. regression values for every proposal
  510. scale_factor_wh (Tensor): tensors of shape [batch_size, 2] the scalor of per img
  511. img_whwh (Tensor): tensors of shape [batch_size, 4]
  512. Returns:
  513. bbox_pred (Tensor): tensors of shape [num_boxes, 6] Each row has 6 values:
  514. [label, confidence, xmin, ymin, xmax, ymax]
  515. bbox_num (Tensor): tensors of shape [batch_size] the number of RoIs in each image.
  516. """
  517. assert len(box_cls) == len(scale_factor_wh) == len(img_whwh)
  518. img_wh = img_whwh[:, :2]
  519. scores = F.sigmoid(box_cls)
  520. labels = paddle.arange(0, self.num_classes). \
  521. unsqueeze(0).tile([self.num_proposals, 1]).flatten(start_axis=0, stop_axis=1)
  522. classes_all = []
  523. scores_all = []
  524. boxes_all = []
  525. for i, (scores_per_image,
  526. box_pred_per_image) in enumerate(zip(scores, box_pred)):
  527. scores_per_image, topk_indices = scores_per_image.flatten(
  528. 0, 1).topk(
  529. self.num_proposals, sorted=False)
  530. labels_per_image = paddle.gather(labels, topk_indices, axis=0)
  531. box_pred_per_image = box_pred_per_image.reshape([-1, 1, 4]).tile(
  532. [1, self.num_classes, 1]).reshape([-1, 4])
  533. box_pred_per_image = paddle.gather(
  534. box_pred_per_image, topk_indices, axis=0)
  535. classes_all.append(labels_per_image)
  536. scores_all.append(scores_per_image)
  537. boxes_all.append(box_pred_per_image)
  538. bbox_num = paddle.zeros([len(scale_factor_wh)], dtype="int32")
  539. boxes_final = []
  540. for i in range(len(scale_factor_wh)):
  541. classes = classes_all[i]
  542. boxes = boxes_all[i]
  543. scores = scores_all[i]
  544. boxes[:, 0::2] = paddle.clip(
  545. boxes[:, 0::2], min=0,
  546. max=img_wh[i][0]) / scale_factor_wh[i][0]
  547. boxes[:, 1::2] = paddle.clip(
  548. boxes[:, 1::2], min=0,
  549. max=img_wh[i][1]) / scale_factor_wh[i][1]
  550. boxes_w, boxes_h = (boxes[:, 2] - boxes[:, 0]).numpy(), (
  551. boxes[:, 3] - boxes[:, 1]).numpy()
  552. keep = (boxes_w > 1.) & (boxes_h > 1.)
  553. if (keep.sum() == 0):
  554. bboxes = paddle.zeros([1, 6]).astype("float32")
  555. else:
  556. boxes = paddle.to_tensor(boxes.numpy()[keep]).astype("float32")
  557. classes = paddle.to_tensor(classes.numpy()[keep]).astype(
  558. "float32").unsqueeze(-1)
  559. scores = paddle.to_tensor(scores.numpy()[keep]).astype(
  560. "float32").unsqueeze(-1)
  561. bboxes = paddle.concat([classes, scores, boxes], axis=-1)
  562. boxes_final.append(bboxes)
  563. bbox_num[i] = bboxes.shape[0]
  564. bbox_pred = paddle.concat(boxes_final)
  565. return bbox_pred, bbox_num