mask_rcnn.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334
  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. from collections import OrderedDict
  18. import copy
  19. import paddle.fluid as fluid
  20. from .fpn import FPN
  21. from .rpn_head import (RPNHead, FPNRPNHead)
  22. from .roi_extractor import (RoIAlign, FPNRoIAlign)
  23. from .bbox_head import (BBoxHead, TwoFCHead)
  24. from .mask_head import MaskHead
  25. from ..resnet import ResNetC5
  26. __all__ = ['MaskRCNN']
  27. class MaskRCNN(object):
  28. """
  29. Mask R-CNN architecture, see https://arxiv.org/abs/1703.06870
  30. Args:
  31. backbone (object): backbone instance
  32. rpn_head (object): `RPNhead` instance
  33. roi_extractor (object): ROI extractor instance
  34. bbox_head (object): `BBoxHead` instance
  35. mask_head (object): `MaskHead` instance
  36. fpn (object): feature pyramid network instance
  37. """
  38. def __init__(
  39. self,
  40. backbone,
  41. num_classes=81,
  42. mode='train',
  43. with_fpn=False,
  44. fpn=None,
  45. num_chan=256,
  46. min_level=2,
  47. max_level=6,
  48. spatial_scale=[1. / 32., 1. / 16., 1. / 8., 1. / 4.],
  49. #rpn_head
  50. rpn_only=False,
  51. rpn_head=None,
  52. anchor_sizes=[32, 64, 128, 256, 512],
  53. aspect_ratios=[0.5, 1.0, 2.0],
  54. rpn_batch_size_per_im=256,
  55. rpn_fg_fraction=0.5,
  56. rpn_positive_overlap=0.7,
  57. rpn_negative_overlap=0.3,
  58. train_pre_nms_top_n=12000,
  59. train_post_nms_top_n=2000,
  60. train_nms_thresh=0.7,
  61. test_pre_nms_top_n=6000,
  62. test_post_nms_top_n=1000,
  63. test_nms_thresh=0.7,
  64. #roi_extractor
  65. roi_extractor=None,
  66. #bbox_head
  67. bbox_head=None,
  68. keep_top_k=100,
  69. nms_threshold=0.5,
  70. score_threshold=0.05,
  71. #MaskHead
  72. mask_head=None,
  73. num_convs=0,
  74. mask_head_resolution=14,
  75. #bbox_assigner
  76. batch_size_per_im=512,
  77. fg_fraction=.25,
  78. fg_thresh=.5,
  79. bg_thresh_hi=.5,
  80. bg_thresh_lo=0.,
  81. bbox_reg_weights=[0.1, 0.1, 0.2, 0.2]):
  82. super(MaskRCNN, self).__init__()
  83. self.backbone = backbone
  84. self.mode = mode
  85. if with_fpn and fpn is None:
  86. fpn = FPN(
  87. num_chan=num_chan,
  88. min_level=min_level,
  89. max_level=max_level,
  90. spatial_scale=spatial_scale)
  91. self.fpn = fpn
  92. self.num_classes = num_classes
  93. if rpn_head is None:
  94. if self.fpn is None:
  95. rpn_head = RPNHead(
  96. anchor_sizes=anchor_sizes,
  97. aspect_ratios=aspect_ratios,
  98. rpn_batch_size_per_im=rpn_batch_size_per_im,
  99. rpn_fg_fraction=rpn_fg_fraction,
  100. rpn_positive_overlap=rpn_positive_overlap,
  101. rpn_negative_overlap=rpn_negative_overlap,
  102. train_pre_nms_top_n=train_pre_nms_top_n,
  103. train_post_nms_top_n=train_post_nms_top_n,
  104. train_nms_thresh=train_nms_thresh,
  105. test_pre_nms_top_n=test_pre_nms_top_n,
  106. test_post_nms_top_n=test_post_nms_top_n,
  107. test_nms_thresh=test_nms_thresh)
  108. else:
  109. rpn_head = FPNRPNHead(
  110. anchor_start_size=anchor_sizes[0],
  111. aspect_ratios=aspect_ratios,
  112. num_chan=self.fpn.num_chan,
  113. min_level=self.fpn.min_level,
  114. max_level=self.fpn.max_level,
  115. rpn_batch_size_per_im=rpn_batch_size_per_im,
  116. rpn_fg_fraction=rpn_fg_fraction,
  117. rpn_positive_overlap=rpn_positive_overlap,
  118. rpn_negative_overlap=rpn_negative_overlap,
  119. train_pre_nms_top_n=train_pre_nms_top_n,
  120. train_post_nms_top_n=train_post_nms_top_n,
  121. train_nms_thresh=train_nms_thresh,
  122. test_pre_nms_top_n=test_pre_nms_top_n,
  123. test_post_nms_top_n=test_post_nms_top_n,
  124. test_nms_thresh=test_nms_thresh)
  125. self.rpn_head = rpn_head
  126. if roi_extractor is None:
  127. if self.fpn is None:
  128. roi_extractor = RoIAlign(
  129. resolution=14,
  130. spatial_scale=1. / 2**self.backbone.feature_maps[0])
  131. else:
  132. roi_extractor = FPNRoIAlign(sampling_ratio=2)
  133. self.roi_extractor = roi_extractor
  134. if bbox_head is None:
  135. if self.fpn is None:
  136. head = ResNetC5(
  137. layers=self.backbone.layers,
  138. norm_type=self.backbone.norm_type,
  139. freeze_norm=self.backbone.freeze_norm)
  140. else:
  141. head = TwoFCHead()
  142. bbox_head = BBoxHead(
  143. head=head,
  144. keep_top_k=keep_top_k,
  145. nms_threshold=nms_threshold,
  146. score_threshold=score_threshold,
  147. num_classes=num_classes)
  148. self.bbox_head = bbox_head
  149. if mask_head is None:
  150. mask_head = MaskHead(
  151. num_convs=num_convs,
  152. resolution=mask_head_resolution,
  153. num_classes=num_classes)
  154. self.mask_head = mask_head
  155. self.batch_size_per_im = batch_size_per_im
  156. self.fg_fraction = fg_fraction
  157. self.fg_thresh = fg_thresh
  158. self.bg_thresh_hi = bg_thresh_hi
  159. self.bg_thresh_lo = bg_thresh_lo
  160. self.bbox_reg_weights = bbox_reg_weights
  161. self.rpn_only = rpn_only
  162. def build_net(self, inputs):
  163. im = inputs['image']
  164. im_info = inputs['im_info']
  165. # backbone
  166. body_feats = self.backbone(im)
  167. # FPN
  168. spatial_scale = None
  169. if self.fpn is not None:
  170. body_feats, spatial_scale = self.fpn.get_output(body_feats)
  171. # RPN proposals
  172. rois = self.rpn_head.get_proposals(body_feats, im_info, mode=self.mode)
  173. if self.mode == 'train':
  174. rpn_loss = self.rpn_head.get_loss(im_info, inputs['gt_box'],
  175. inputs['is_crowd'])
  176. outputs = fluid.layers.generate_proposal_labels(
  177. rpn_rois=rois,
  178. gt_classes=inputs['gt_label'],
  179. is_crowd=inputs['is_crowd'],
  180. gt_boxes=inputs['gt_box'],
  181. im_info=inputs['im_info'],
  182. batch_size_per_im=self.batch_size_per_im,
  183. fg_fraction=self.fg_fraction,
  184. fg_thresh=self.fg_thresh,
  185. bg_thresh_hi=self.bg_thresh_hi,
  186. bg_thresh_lo=self.bg_thresh_lo,
  187. bbox_reg_weights=self.bbox_reg_weights,
  188. class_nums=self.num_classes,
  189. use_random=self.rpn_head.use_random)
  190. rois = outputs[0]
  191. labels_int32 = outputs[1]
  192. if self.fpn is None:
  193. last_feat = body_feats[list(body_feats.keys())[-1]]
  194. roi_feat = self.roi_extractor(last_feat, rois)
  195. else:
  196. roi_feat = self.roi_extractor(body_feats, rois, spatial_scale)
  197. loss = self.bbox_head.get_loss(roi_feat, labels_int32,
  198. *outputs[2:])
  199. loss.update(rpn_loss)
  200. mask_rois, roi_has_mask_int32, mask_int32 = fluid.layers.generate_mask_labels(
  201. rois=rois,
  202. gt_classes=inputs['gt_label'],
  203. is_crowd=inputs['is_crowd'],
  204. gt_segms=inputs['gt_mask'],
  205. im_info=inputs['im_info'],
  206. labels_int32=labels_int32,
  207. num_classes=self.num_classes,
  208. resolution=self.mask_head.resolution)
  209. if self.fpn is None:
  210. bbox_head_feat = self.bbox_head.get_head_feat()
  211. feat = fluid.layers.gather(bbox_head_feat, roi_has_mask_int32)
  212. else:
  213. feat = self.roi_extractor(
  214. body_feats, mask_rois, spatial_scale, is_mask=True)
  215. mask_loss = self.mask_head.get_loss(feat, mask_int32)
  216. loss.update(mask_loss)
  217. total_loss = fluid.layers.sum(list(loss.values()))
  218. loss.update({'loss': total_loss})
  219. return loss
  220. else:
  221. if self.rpn_only:
  222. im_scale = fluid.layers.slice(
  223. im_info, [1], starts=[2], ends=[3])
  224. im_scale = fluid.layers.sequence_expand(im_scale, rois)
  225. rois = rois / im_scale
  226. return {'proposal': rois}
  227. mask_name = 'mask_pred'
  228. mask_pred, bbox_pred = self._eval(body_feats, mask_name, rois,
  229. im_info, inputs['im_shape'],
  230. spatial_scale)
  231. return OrderedDict(zip(['bbox', 'mask'], [bbox_pred, mask_pred]))
  232. def _eval(self,
  233. body_feats,
  234. mask_name,
  235. rois,
  236. im_info,
  237. im_shape,
  238. spatial_scale,
  239. bbox_pred=None):
  240. if not bbox_pred:
  241. if self.fpn is None:
  242. last_feat = body_feats[list(body_feats.keys())[-1]]
  243. roi_feat = self.roi_extractor(last_feat, rois)
  244. else:
  245. roi_feat = self.roi_extractor(body_feats, rois, spatial_scale)
  246. bbox_pred = self.bbox_head.get_prediction(roi_feat, rois, im_info,
  247. im_shape)
  248. bbox_pred = bbox_pred['bbox']
  249. # share weight
  250. bbox_shape = fluid.layers.shape(bbox_pred)
  251. bbox_size = fluid.layers.reduce_prod(bbox_shape)
  252. bbox_size = fluid.layers.reshape(bbox_size, [1, 1])
  253. size = fluid.layers.fill_constant([1, 1], value=6, dtype='int32')
  254. cond = fluid.layers.less_than(x=bbox_size, y=size)
  255. mask_pred = fluid.layers.create_global_var(
  256. shape=[1],
  257. value=0.0,
  258. dtype='float32',
  259. persistable=False,
  260. name=mask_name)
  261. with fluid.layers.control_flow.Switch() as switch:
  262. with switch.case(cond):
  263. fluid.layers.assign(input=bbox_pred, output=mask_pred)
  264. with switch.default():
  265. bbox = fluid.layers.slice(bbox_pred, [1], starts=[2], ends=[6])
  266. im_scale = fluid.layers.slice(
  267. im_info, [1], starts=[2], ends=[3])
  268. im_scale = fluid.layers.sequence_expand(im_scale, bbox)
  269. mask_rois = bbox * im_scale
  270. if self.fpn is None:
  271. last_feat = body_feats[list(body_feats.keys())[-1]]
  272. mask_feat = self.roi_extractor(last_feat, mask_rois)
  273. mask_feat = self.bbox_head.get_head_feat(mask_feat)
  274. else:
  275. mask_feat = self.roi_extractor(
  276. body_feats, mask_rois, spatial_scale, is_mask=True)
  277. mask_out = self.mask_head.get_prediction(mask_feat, bbox)
  278. fluid.layers.assign(input=mask_out, output=mask_pred)
  279. return mask_pred, bbox_pred
  280. def generate_inputs(self):
  281. inputs = OrderedDict()
  282. inputs['image'] = fluid.data(
  283. dtype='float32', shape=[None, 3, None, None], name='image')
  284. if self.mode == 'train':
  285. inputs['im_info'] = fluid.data(
  286. dtype='float32', shape=[None, 3], name='im_info')
  287. inputs['gt_box'] = fluid.data(
  288. dtype='float32', shape=[None, 4], lod_level=1, name='gt_box')
  289. inputs['gt_label'] = fluid.data(
  290. dtype='int32', shape=[None, 1], lod_level=1, name='gt_label')
  291. inputs['is_crowd'] = fluid.data(
  292. dtype='int32', shape=[None, 1], lod_level=1, name='is_crowd')
  293. inputs['gt_mask'] = fluid.data(
  294. dtype='float32', shape=[None, 2], lod_level=3, name='gt_mask')
  295. elif self.mode == 'eval':
  296. inputs['im_info'] = fluid.data(
  297. dtype='float32', shape=[None, 3], name='im_info')
  298. inputs['im_id'] = fluid.data(
  299. dtype='int64', shape=[None, 1], name='im_id')
  300. inputs['im_shape'] = fluid.data(
  301. dtype='float32', shape=[None, 3], name='im_shape')
  302. elif self.mode == 'test':
  303. inputs['im_info'] = fluid.data(
  304. dtype='float32', shape=[None, 3], name='im_info')
  305. inputs['im_shape'] = fluid.data(
  306. dtype='float32', shape=[None, 3], name='im_shape')
  307. return inputs