mask_rcnn.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353
  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, HRFPN)
  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. fixed_input_shape=None,
  83. input_channel=3):
  84. super(MaskRCNN, self).__init__()
  85. self.backbone = backbone
  86. self.mode = mode
  87. if with_fpn and fpn is None:
  88. if self.backbone.__class__.__name__.startswith('HRNet'):
  89. fpn = HRFPN()
  90. fpn.min_level = 2
  91. fpn.max_level = 6
  92. else:
  93. fpn = FPN(num_chan=num_chan,
  94. min_level=min_level,
  95. max_level=max_level,
  96. spatial_scale=spatial_scale)
  97. self.fpn = fpn
  98. self.num_classes = num_classes
  99. if rpn_head is None:
  100. if self.fpn is None:
  101. rpn_head = RPNHead(
  102. anchor_sizes=anchor_sizes,
  103. aspect_ratios=aspect_ratios,
  104. rpn_batch_size_per_im=rpn_batch_size_per_im,
  105. rpn_fg_fraction=rpn_fg_fraction,
  106. rpn_positive_overlap=rpn_positive_overlap,
  107. rpn_negative_overlap=rpn_negative_overlap,
  108. train_pre_nms_top_n=train_pre_nms_top_n,
  109. train_post_nms_top_n=train_post_nms_top_n,
  110. train_nms_thresh=train_nms_thresh,
  111. test_pre_nms_top_n=test_pre_nms_top_n,
  112. test_post_nms_top_n=test_post_nms_top_n,
  113. test_nms_thresh=test_nms_thresh)
  114. else:
  115. rpn_head = FPNRPNHead(
  116. anchor_start_size=anchor_sizes[0],
  117. aspect_ratios=aspect_ratios,
  118. num_chan=self.fpn.num_chan,
  119. min_level=self.fpn.min_level,
  120. max_level=self.fpn.max_level,
  121. rpn_batch_size_per_im=rpn_batch_size_per_im,
  122. rpn_fg_fraction=rpn_fg_fraction,
  123. rpn_positive_overlap=rpn_positive_overlap,
  124. rpn_negative_overlap=rpn_negative_overlap,
  125. train_pre_nms_top_n=train_pre_nms_top_n,
  126. train_post_nms_top_n=train_post_nms_top_n,
  127. train_nms_thresh=train_nms_thresh,
  128. test_pre_nms_top_n=test_pre_nms_top_n,
  129. test_post_nms_top_n=test_post_nms_top_n,
  130. test_nms_thresh=test_nms_thresh)
  131. self.rpn_head = rpn_head
  132. if roi_extractor is None:
  133. if self.fpn is None:
  134. roi_extractor = RoIAlign(
  135. resolution=14,
  136. spatial_scale=1. / 2**self.backbone.feature_maps[0])
  137. else:
  138. roi_extractor = FPNRoIAlign(sampling_ratio=2)
  139. self.roi_extractor = roi_extractor
  140. if bbox_head is None:
  141. if self.fpn is None:
  142. head = ResNetC5(
  143. layers=self.backbone.layers,
  144. norm_type=self.backbone.norm_type,
  145. freeze_norm=self.backbone.freeze_norm)
  146. else:
  147. head = TwoFCHead()
  148. bbox_head = BBoxHead(
  149. head=head,
  150. keep_top_k=keep_top_k,
  151. nms_threshold=nms_threshold,
  152. score_threshold=score_threshold,
  153. num_classes=num_classes)
  154. self.bbox_head = bbox_head
  155. if mask_head is None:
  156. mask_head = MaskHead(
  157. num_convs=num_convs,
  158. resolution=mask_head_resolution,
  159. num_classes=num_classes)
  160. self.mask_head = mask_head
  161. self.batch_size_per_im = batch_size_per_im
  162. self.fg_fraction = fg_fraction
  163. self.fg_thresh = fg_thresh
  164. self.bg_thresh_hi = bg_thresh_hi
  165. self.bg_thresh_lo = bg_thresh_lo
  166. self.bbox_reg_weights = bbox_reg_weights
  167. self.rpn_only = rpn_only
  168. self.fixed_input_shape = fixed_input_shape
  169. self.input_channel = input_channel
  170. def build_net(self, inputs):
  171. im = inputs['image']
  172. im_info = inputs['im_info']
  173. # backbone
  174. body_feats = self.backbone(im)
  175. # FPN
  176. spatial_scale = None
  177. if self.fpn is not None:
  178. body_feats, spatial_scale = self.fpn.get_output(body_feats)
  179. # RPN proposals
  180. rois = self.rpn_head.get_proposals(body_feats, im_info, mode=self.mode)
  181. if self.mode == 'train':
  182. rpn_loss = self.rpn_head.get_loss(im_info, inputs['gt_box'],
  183. inputs['is_crowd'])
  184. outputs = fluid.layers.generate_proposal_labels(
  185. rpn_rois=rois,
  186. gt_classes=inputs['gt_label'],
  187. is_crowd=inputs['is_crowd'],
  188. gt_boxes=inputs['gt_box'],
  189. im_info=inputs['im_info'],
  190. batch_size_per_im=self.batch_size_per_im,
  191. fg_fraction=self.fg_fraction,
  192. fg_thresh=self.fg_thresh,
  193. bg_thresh_hi=self.bg_thresh_hi,
  194. bg_thresh_lo=self.bg_thresh_lo,
  195. bbox_reg_weights=self.bbox_reg_weights,
  196. class_nums=self.num_classes,
  197. use_random=self.rpn_head.use_random)
  198. rois = outputs[0]
  199. labels_int32 = outputs[1]
  200. if self.fpn is None:
  201. last_feat = body_feats[list(body_feats.keys())[-1]]
  202. roi_feat = self.roi_extractor(last_feat, rois)
  203. else:
  204. roi_feat = self.roi_extractor(body_feats, rois, spatial_scale)
  205. loss = self.bbox_head.get_loss(roi_feat, labels_int32,
  206. *outputs[2:])
  207. loss.update(rpn_loss)
  208. mask_rois, roi_has_mask_int32, mask_int32 = fluid.layers.generate_mask_labels(
  209. rois=rois,
  210. gt_classes=inputs['gt_label'],
  211. is_crowd=inputs['is_crowd'],
  212. gt_segms=inputs['gt_mask'],
  213. im_info=inputs['im_info'],
  214. labels_int32=labels_int32,
  215. num_classes=self.num_classes,
  216. resolution=self.mask_head.resolution)
  217. if self.fpn is None:
  218. bbox_head_feat = self.bbox_head.get_head_feat()
  219. feat = fluid.layers.gather(bbox_head_feat, roi_has_mask_int32)
  220. else:
  221. feat = self.roi_extractor(
  222. body_feats, mask_rois, spatial_scale, is_mask=True)
  223. mask_loss = self.mask_head.get_loss(feat, mask_int32)
  224. loss.update(mask_loss)
  225. total_loss = fluid.layers.sum(list(loss.values()))
  226. loss.update({'loss': total_loss})
  227. return loss
  228. else:
  229. if self.rpn_only:
  230. im_scale = fluid.layers.slice(
  231. im_info, [1], starts=[2], ends=[3])
  232. im_scale = fluid.layers.sequence_expand(im_scale, rois)
  233. rois = rois / im_scale
  234. return {'proposal': rois}
  235. mask_name = 'mask_pred'
  236. mask_pred, bbox_pred = self._eval(body_feats, mask_name, rois,
  237. im_info, inputs['im_shape'],
  238. spatial_scale)
  239. return OrderedDict(zip(['bbox', 'mask'], [bbox_pred, mask_pred]))
  240. def _eval(self,
  241. body_feats,
  242. mask_name,
  243. rois,
  244. im_info,
  245. im_shape,
  246. spatial_scale,
  247. bbox_pred=None):
  248. if not bbox_pred:
  249. if self.fpn is None:
  250. last_feat = body_feats[list(body_feats.keys())[-1]]
  251. roi_feat = self.roi_extractor(last_feat, rois)
  252. else:
  253. roi_feat = self.roi_extractor(body_feats, rois, spatial_scale)
  254. bbox_pred = self.bbox_head.get_prediction(roi_feat, rois, im_info,
  255. im_shape)
  256. bbox_pred = bbox_pred['bbox']
  257. # share weight
  258. bbox_shape = fluid.layers.shape(bbox_pred)
  259. bbox_size = fluid.layers.reduce_prod(bbox_shape)
  260. bbox_size = fluid.layers.reshape(bbox_size, [1, 1])
  261. size = fluid.layers.fill_constant([1, 1], value=6, dtype='int32')
  262. cond = fluid.layers.less_than(x=bbox_size, y=size)
  263. mask_pred = fluid.layers.create_global_var(
  264. shape=[1],
  265. value=0.0,
  266. dtype='float32',
  267. persistable=False,
  268. name=mask_name)
  269. with fluid.layers.control_flow.Switch() as switch:
  270. with switch.case(cond):
  271. fluid.layers.assign(input=bbox_pred, output=mask_pred)
  272. with switch.default():
  273. bbox = fluid.layers.slice(bbox_pred, [1], starts=[2], ends=[6])
  274. im_scale = fluid.layers.slice(
  275. im_info, [1], starts=[2], ends=[3])
  276. im_scale = fluid.layers.sequence_expand(im_scale, bbox)
  277. mask_rois = bbox * im_scale
  278. if self.fpn is None:
  279. last_feat = body_feats[list(body_feats.keys())[-1]]
  280. mask_feat = self.roi_extractor(last_feat, mask_rois)
  281. mask_feat = self.bbox_head.get_head_feat(mask_feat)
  282. else:
  283. mask_feat = self.roi_extractor(
  284. body_feats, mask_rois, spatial_scale, is_mask=True)
  285. mask_out = self.mask_head.get_prediction(mask_feat, bbox)
  286. fluid.layers.assign(input=mask_out, output=mask_pred)
  287. return mask_pred, bbox_pred
  288. def generate_inputs(self):
  289. inputs = OrderedDict()
  290. if self.fixed_input_shape is not None:
  291. input_shape = [
  292. None, self.input_channel, self.fixed_input_shape[1],
  293. self.fixed_input_shape[0]
  294. ]
  295. inputs['image'] = fluid.data(
  296. dtype='float32', shape=input_shape, name='image')
  297. else:
  298. inputs['image'] = fluid.data(
  299. dtype='float32',
  300. shape=[None, self.input_channel, None, None],
  301. name='image')
  302. if self.mode == 'train':
  303. inputs['im_info'] = fluid.data(
  304. dtype='float32', shape=[None, 3], name='im_info')
  305. inputs['gt_box'] = fluid.data(
  306. dtype='float32', shape=[None, 4], lod_level=1, name='gt_box')
  307. inputs['gt_label'] = fluid.data(
  308. dtype='int32', shape=[None, 1], lod_level=1, name='gt_label')
  309. inputs['is_crowd'] = fluid.data(
  310. dtype='int32', shape=[None, 1], lod_level=1, name='is_crowd')
  311. inputs['gt_mask'] = fluid.data(
  312. dtype='float32', shape=[None, 2], lod_level=3, name='gt_mask')
  313. elif self.mode == 'eval':
  314. inputs['im_info'] = fluid.data(
  315. dtype='float32', shape=[None, 3], name='im_info')
  316. inputs['im_id'] = fluid.data(
  317. dtype='int64', shape=[None, 1], name='im_id')
  318. inputs['im_shape'] = fluid.data(
  319. dtype='float32', shape=[None, 3], name='im_shape')
  320. elif self.mode == 'test':
  321. inputs['im_info'] = fluid.data(
  322. dtype='float32', shape=[None, 3], name='im_info')
  323. inputs['im_shape'] = fluid.data(
  324. dtype='float32', shape=[None, 3], name='im_shape')
  325. return inputs