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