fcos.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  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. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import paddle
  18. from ppdet.core.workspace import register, create
  19. from .meta_arch import BaseArch
  20. __all__ = ['FCOS']
  21. @register
  22. class FCOS(BaseArch):
  23. """
  24. FCOS network, see https://arxiv.org/abs/1904.01355
  25. Args:
  26. backbone (object): backbone instance
  27. neck (object): 'FPN' instance
  28. fcos_head (object): 'FCOSHead' instance
  29. post_process (object): 'FCOSPostProcess' instance
  30. """
  31. __category__ = 'architecture'
  32. __inject__ = ['fcos_post_process']
  33. def __init__(self,
  34. backbone,
  35. neck,
  36. fcos_head='FCOSHead',
  37. fcos_post_process='FCOSPostProcess'):
  38. super(FCOS, self).__init__()
  39. self.backbone = backbone
  40. self.neck = neck
  41. self.fcos_head = fcos_head
  42. self.fcos_post_process = fcos_post_process
  43. @classmethod
  44. def from_config(cls, cfg, *args, **kwargs):
  45. backbone = create(cfg['backbone'])
  46. kwargs = {'input_shape': backbone.out_shape}
  47. neck = create(cfg['neck'], **kwargs)
  48. kwargs = {'input_shape': neck.out_shape}
  49. fcos_head = create(cfg['fcos_head'], **kwargs)
  50. return {
  51. 'backbone': backbone,
  52. 'neck': neck,
  53. "fcos_head": fcos_head,
  54. }
  55. def _forward(self):
  56. body_feats = self.backbone(self.inputs)
  57. fpn_feats = self.neck(body_feats)
  58. fcos_head_outs = self.fcos_head(fpn_feats, self.training)
  59. if not self.training:
  60. scale_factor = self.inputs['scale_factor']
  61. bboxes = self.fcos_post_process(fcos_head_outs, scale_factor)
  62. return bboxes
  63. else:
  64. return fcos_head_outs
  65. def get_loss(self, ):
  66. loss = {}
  67. tag_labels, tag_bboxes, tag_centerness = [], [], []
  68. for i in range(len(self.fcos_head.fpn_stride)):
  69. # labels, reg_target, centerness
  70. k_lbl = 'labels{}'.format(i)
  71. if k_lbl in self.inputs:
  72. tag_labels.append(self.inputs[k_lbl])
  73. k_box = 'reg_target{}'.format(i)
  74. if k_box in self.inputs:
  75. tag_bboxes.append(self.inputs[k_box])
  76. k_ctn = 'centerness{}'.format(i)
  77. if k_ctn in self.inputs:
  78. tag_centerness.append(self.inputs[k_ctn])
  79. fcos_head_outs = self._forward()
  80. loss_fcos = self.fcos_head.get_loss(fcos_head_outs, tag_labels,
  81. tag_bboxes, tag_centerness)
  82. loss.update(loss_fcos)
  83. total_loss = paddle.add_n(list(loss.values()))
  84. loss.update({'loss': total_loss})
  85. return loss
  86. def get_pred(self):
  87. bbox_pred, bbox_num = self._forward()
  88. output = {'bbox': bbox_pred, 'bbox_num': bbox_num}
  89. return output