anchor_generator.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  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. # The code is based on
  15. # https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/anchor_generator.py
  16. import math
  17. import paddle
  18. import paddle.nn as nn
  19. from paddlex.ppdet.core.workspace import register
  20. @register
  21. class AnchorGenerator(nn.Layer):
  22. """
  23. Generate anchors according to the feature maps
  24. Args:
  25. anchor_sizes (list[float] | list[list[float]]): The anchor sizes at
  26. each feature point. list[float] means all feature levels share the
  27. same sizes. list[list[float]] means the anchor sizes for
  28. each level. The sizes stand for the scale of input size.
  29. aspect_ratios (list[float] | list[list[float]]): The aspect ratios at
  30. each feature point. list[float] means all feature levels share the
  31. same ratios. list[list[float]] means the aspect ratios for
  32. each level.
  33. strides (list[float]): The strides of feature maps which generate
  34. anchors
  35. offset (float): The offset of the coordinate of anchors, default 0.
  36. """
  37. def __init__(self,
  38. anchor_sizes=[32, 64, 128, 256, 512],
  39. aspect_ratios=[0.5, 1.0, 2.0],
  40. strides=[16.0],
  41. variance=[1.0, 1.0, 1.0, 1.0],
  42. offset=0.):
  43. super(AnchorGenerator, self).__init__()
  44. self.anchor_sizes = anchor_sizes
  45. self.aspect_ratios = aspect_ratios
  46. self.strides = strides
  47. self.variance = variance
  48. self.cell_anchors = self._calculate_anchors(len(strides))
  49. self.offset = offset
  50. def _broadcast_params(self, params, num_features):
  51. if not isinstance(params[0], (list, tuple)): # list[float]
  52. return [params] * num_features
  53. if len(params) == 1:
  54. return list(params) * num_features
  55. return params
  56. def generate_cell_anchors(self, sizes, aspect_ratios):
  57. anchors = []
  58. for size in sizes:
  59. area = size**2.0
  60. for aspect_ratio in aspect_ratios:
  61. w = math.sqrt(area / aspect_ratio)
  62. h = aspect_ratio * w
  63. x0, y0, x1, y1 = -w / 2.0, -h / 2.0, w / 2.0, h / 2.0
  64. anchors.append([x0, y0, x1, y1])
  65. return paddle.to_tensor(anchors, dtype='float32')
  66. def _calculate_anchors(self, num_features):
  67. sizes = self._broadcast_params(self.anchor_sizes, num_features)
  68. aspect_ratios = self._broadcast_params(self.aspect_ratios,
  69. num_features)
  70. cell_anchors = [
  71. self.generate_cell_anchors(s, a)
  72. for s, a in zip(sizes, aspect_ratios)
  73. ]
  74. [
  75. self.register_buffer(
  76. t.name, t, persistable=False) for t in cell_anchors
  77. ]
  78. return cell_anchors
  79. def _create_grid_offsets(self, size, stride, offset):
  80. grid_height, grid_width = size[0], size[1]
  81. shifts_x = paddle.arange(
  82. offset * stride, grid_width * stride, step=stride, dtype='float32')
  83. shifts_y = paddle.arange(
  84. offset * stride,
  85. grid_height * stride,
  86. step=stride,
  87. dtype='float32')
  88. shift_y, shift_x = paddle.meshgrid(shifts_y, shifts_x)
  89. shift_x = paddle.reshape(shift_x, [-1])
  90. shift_y = paddle.reshape(shift_y, [-1])
  91. return shift_x, shift_y
  92. def _grid_anchors(self, grid_sizes):
  93. anchors = []
  94. for size, stride, base_anchors in zip(grid_sizes, self.strides,
  95. self.cell_anchors):
  96. shift_x, shift_y = self._create_grid_offsets(size, stride,
  97. self.offset)
  98. shifts = paddle.stack((shift_x, shift_y, shift_x, shift_y), axis=1)
  99. shifts = paddle.reshape(shifts, [-1, 1, 4])
  100. base_anchors = paddle.reshape(base_anchors, [1, -1, 4])
  101. anchors.append(paddle.reshape(shifts + base_anchors, [-1, 4]))
  102. return anchors
  103. def forward(self, input):
  104. grid_sizes = [paddle.shape(feature_map)[-2:] for feature_map in input]
  105. anchors_over_all_feature_maps = self._grid_anchors(grid_sizes)
  106. return anchors_over_all_feature_maps
  107. @property
  108. def num_anchors(self):
  109. """
  110. Returns:
  111. int: number of anchors at every pixel
  112. location, on that feature map.
  113. For example, if at every pixel we use anchors of 3 aspect
  114. ratios and 5 sizes, the number of anchors is 15.
  115. For FPN models, `num_anchors` on every feature map is the same.
  116. """
  117. return len(self.cell_anchors[0])