keypoint_hrnet.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203
  1. # Copyright (c) 2021 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. import numpy as np
  19. import math
  20. from paddlex.ppdet.core.workspace import register, create
  21. from .meta_arch import BaseArch
  22. from ..keypoint_utils import transform_preds
  23. from .. import layers as L
  24. __all__ = ['TopDownHRNet']
  25. @register
  26. class TopDownHRNet(BaseArch):
  27. __category__ = 'architecture'
  28. __inject__ = ['loss']
  29. def __init__(self,
  30. width,
  31. num_joints,
  32. backbone='HRNet',
  33. loss='KeyPointMSELoss',
  34. post_process='HRNetPostProcess',
  35. flip_perm=None,
  36. flip=True,
  37. shift_heatmap=True):
  38. """
  39. HRNnet network, see https://arxiv.org/abs/1902.09212
  40. Args:
  41. backbone (nn.Layer): backbone instance
  42. post_process (object): `HRNetPostProcess` instance
  43. flip_perm (list): The left-right joints exchange order list
  44. """
  45. super(TopDownHRNet, self).__init__()
  46. self.backbone = backbone
  47. self.post_process = HRNetPostProcess()
  48. self.loss = loss
  49. self.flip_perm = flip_perm
  50. self.flip = flip
  51. self.final_conv = L.Conv2d(width, num_joints, 1, 1, 0, bias=True)
  52. self.shift_heatmap = shift_heatmap
  53. self.deploy = False
  54. @classmethod
  55. def from_config(cls, cfg, *args, **kwargs):
  56. # backbone
  57. backbone = create(cfg['backbone'])
  58. return {'backbone': backbone, }
  59. def _forward(self):
  60. feats = self.backbone(self.inputs)
  61. hrnet_outputs = self.final_conv(feats[0])
  62. if self.training:
  63. return self.loss(hrnet_outputs, self.inputs)
  64. elif self.deploy:
  65. return hrnet_outputs
  66. else:
  67. if self.flip:
  68. self.inputs['image'] = self.inputs['image'].flip([3])
  69. feats = self.backbone(self.inputs)
  70. output_flipped = self.final_conv(feats[0])
  71. output_flipped = self.flip_back(output_flipped.numpy(),
  72. self.flip_perm)
  73. output_flipped = paddle.to_tensor(output_flipped.copy())
  74. if self.shift_heatmap:
  75. output_flipped[:, :, :, 1:] = output_flipped.clone(
  76. )[:, :, :, 0:-1]
  77. hrnet_outputs = (hrnet_outputs + output_flipped) * 0.5
  78. imshape = (self.inputs['im_shape'].numpy()
  79. )[:, ::-1] if 'im_shape' in self.inputs else None
  80. center = self.inputs['center'].numpy(
  81. ) if 'center' in self.inputs else np.round(imshape / 2.)
  82. scale = self.inputs['scale'].numpy(
  83. ) if 'scale' in self.inputs else imshape / 200.
  84. outputs = self.post_process(hrnet_outputs, center, scale)
  85. return outputs
  86. def get_loss(self):
  87. return self._forward()
  88. def get_pred(self):
  89. res_lst = self._forward()
  90. outputs = {'keypoint': res_lst}
  91. return outputs
  92. def flip_back(self, output_flipped, matched_parts):
  93. assert output_flipped.ndim == 4,\
  94. 'output_flipped should be [batch_size, num_joints, height, width]'
  95. output_flipped = output_flipped[:, :, :, ::-1]
  96. for pair in matched_parts:
  97. tmp = output_flipped[:, pair[0], :, :].copy()
  98. output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
  99. output_flipped[:, pair[1], :, :] = tmp
  100. return output_flipped
  101. class HRNetPostProcess(object):
  102. def get_max_preds(self, heatmaps):
  103. '''get predictions from score maps
  104. Args:
  105. heatmaps: numpy.ndarray([batch_size, num_joints, height, width])
  106. Returns:
  107. preds: numpy.ndarray([batch_size, num_joints, 2]), keypoints coords
  108. maxvals: numpy.ndarray([batch_size, num_joints, 2]), the maximum confidence of the keypoints
  109. '''
  110. assert isinstance(heatmaps,
  111. np.ndarray), 'heatmaps should be numpy.ndarray'
  112. assert heatmaps.ndim == 4, 'batch_images should be 4-ndim'
  113. batch_size = heatmaps.shape[0]
  114. num_joints = heatmaps.shape[1]
  115. width = heatmaps.shape[3]
  116. heatmaps_reshaped = heatmaps.reshape((batch_size, num_joints, -1))
  117. idx = np.argmax(heatmaps_reshaped, 2)
  118. maxvals = np.amax(heatmaps_reshaped, 2)
  119. maxvals = maxvals.reshape((batch_size, num_joints, 1))
  120. idx = idx.reshape((batch_size, num_joints, 1))
  121. preds = np.tile(idx, (1, 1, 2)).astype(np.float32)
  122. preds[:, :, 0] = (preds[:, :, 0]) % width
  123. preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)
  124. pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
  125. pred_mask = pred_mask.astype(np.float32)
  126. preds *= pred_mask
  127. return preds, maxvals
  128. def get_final_preds(self, heatmaps, center, scale):
  129. """the highest heatvalue location with a quarter offset in the
  130. direction from the highest response to the second highest response.
  131. Args:
  132. heatmaps (numpy.ndarray): The predicted heatmaps
  133. center (numpy.ndarray): The boxes center
  134. scale (numpy.ndarray): The scale factor
  135. Returns:
  136. preds: numpy.ndarray([batch_size, num_joints, 2]), keypoints coords
  137. maxvals: numpy.ndarray([batch_size, num_joints, 1]), the maximum confidence of the keypoints
  138. """
  139. coords, maxvals = self.get_max_preds(heatmaps)
  140. heatmap_height = heatmaps.shape[2]
  141. heatmap_width = heatmaps.shape[3]
  142. for n in range(coords.shape[0]):
  143. for p in range(coords.shape[1]):
  144. hm = heatmaps[n][p]
  145. px = int(math.floor(coords[n][p][0] + 0.5))
  146. py = int(math.floor(coords[n][p][1] + 0.5))
  147. if 1 < px < heatmap_width - 1 and 1 < py < heatmap_height - 1:
  148. diff = np.array([
  149. hm[py][px + 1] - hm[py][px - 1],
  150. hm[py + 1][px] - hm[py - 1][px]
  151. ])
  152. coords[n][p] += np.sign(diff) * .25
  153. preds = coords.copy()
  154. # Transform back
  155. for i in range(coords.shape[0]):
  156. preds[i] = transform_preds(coords[i], center[i], scale[i],
  157. [heatmap_width, heatmap_height])
  158. return preds, maxvals
  159. def __call__(self, output, center, scale):
  160. preds, maxvals = self.get_final_preds(output.numpy(), center, scale)
  161. outputs = [[
  162. np.concatenate(
  163. (preds, maxvals), axis=-1), np.mean(
  164. maxvals, axis=1)
  165. ]]
  166. return outputs