keypoint_loss.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228
  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. from itertools import cycle, islice
  18. from collections import abc
  19. import paddle
  20. import paddle.nn as nn
  21. from paddlex.ppdet.core.workspace import register, serializable
  22. __all__ = ['HrHRNetLoss', 'KeyPointMSELoss']
  23. @register
  24. @serializable
  25. class KeyPointMSELoss(nn.Layer):
  26. def __init__(self, use_target_weight=True):
  27. """
  28. KeyPointMSELoss layer
  29. Args:
  30. use_target_weight (bool): whether to use target weight
  31. """
  32. super(KeyPointMSELoss, self).__init__()
  33. self.criterion = nn.MSELoss(reduction='mean')
  34. self.use_target_weight = use_target_weight
  35. def forward(self, output, records):
  36. target = records['target']
  37. target_weight = records['target_weight']
  38. batch_size = output.shape[0]
  39. num_joints = output.shape[1]
  40. heatmaps_pred = output.reshape(
  41. (batch_size, num_joints, -1)).split(num_joints, 1)
  42. heatmaps_gt = target.reshape(
  43. (batch_size, num_joints, -1)).split(num_joints, 1)
  44. loss = 0
  45. for idx in range(num_joints):
  46. heatmap_pred = heatmaps_pred[idx].squeeze()
  47. heatmap_gt = heatmaps_gt[idx].squeeze()
  48. if self.use_target_weight:
  49. loss += 0.5 * self.criterion(
  50. heatmap_pred.multiply(target_weight[:, idx]),
  51. heatmap_gt.multiply(target_weight[:, idx]))
  52. else:
  53. loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)
  54. keypoint_losses = dict()
  55. keypoint_losses['loss'] = loss / num_joints
  56. return keypoint_losses
  57. @register
  58. @serializable
  59. class HrHRNetLoss(nn.Layer):
  60. def __init__(self, num_joints, swahr):
  61. """
  62. HrHRNetLoss layer
  63. Args:
  64. num_joints (int): number of keypoints
  65. """
  66. super(HrHRNetLoss, self).__init__()
  67. if swahr:
  68. self.heatmaploss = HeatMapSWAHRLoss(num_joints)
  69. else:
  70. self.heatmaploss = HeatMapLoss()
  71. self.aeloss = AELoss()
  72. self.ziploss = ZipLoss(
  73. [self.heatmaploss, self.heatmaploss, self.aeloss])
  74. def forward(self, inputs, records):
  75. targets = []
  76. targets.append([records['heatmap_gt1x'], records['mask_1x']])
  77. targets.append([records['heatmap_gt2x'], records['mask_2x']])
  78. targets.append(records['tagmap'])
  79. keypoint_losses = dict()
  80. loss = self.ziploss(inputs, targets)
  81. keypoint_losses['heatmap_loss'] = loss[0] + loss[1]
  82. keypoint_losses['pull_loss'] = loss[2][0]
  83. keypoint_losses['push_loss'] = loss[2][1]
  84. keypoint_losses['loss'] = recursive_sum(loss)
  85. return keypoint_losses
  86. class HeatMapLoss(object):
  87. def __init__(self, loss_factor=1.0):
  88. super(HeatMapLoss, self).__init__()
  89. self.loss_factor = loss_factor
  90. def __call__(self, preds, targets):
  91. heatmap, mask = targets
  92. loss = ((preds - heatmap)**2 * mask.cast('float').unsqueeze(1))
  93. loss = paddle.clip(loss, min=0, max=2).mean()
  94. loss *= self.loss_factor
  95. return loss
  96. class HeatMapSWAHRLoss(object):
  97. def __init__(self, num_joints, loss_factor=1.0):
  98. super(HeatMapSWAHRLoss, self).__init__()
  99. self.loss_factor = loss_factor
  100. self.num_joints = num_joints
  101. def __call__(self, preds, targets):
  102. heatmaps_gt, mask = targets
  103. heatmaps_pred = preds[0]
  104. scalemaps_pred = preds[1]
  105. heatmaps_scaled_gt = paddle.where(
  106. heatmaps_gt > 0, 0.5 * heatmaps_gt *
  107. (1 + (1 +
  108. (scalemaps_pred - 1.) * paddle.log(heatmaps_gt + 1e-10))**2),
  109. heatmaps_gt)
  110. regularizer_loss = paddle.mean(
  111. paddle.pow((scalemaps_pred - 1.) * (heatmaps_gt > 0).astype(float),
  112. 2))
  113. omiga = 0.01
  114. # thres = 2**(-1/omiga), threshold for positive weight
  115. hm_weight = heatmaps_scaled_gt**(
  116. omiga
  117. ) * paddle.abs(1 - heatmaps_pred) + paddle.abs(heatmaps_pred) * (
  118. 1 - heatmaps_scaled_gt**(omiga))
  119. loss = (((heatmaps_pred - heatmaps_scaled_gt)**2) *
  120. mask.cast('float').unsqueeze(1)) * hm_weight
  121. loss = loss.mean()
  122. loss = self.loss_factor * (loss + 1.0 * regularizer_loss)
  123. return loss
  124. class AELoss(object):
  125. def __init__(self, pull_factor=0.001, push_factor=0.001):
  126. super(AELoss, self).__init__()
  127. self.pull_factor = pull_factor
  128. self.push_factor = push_factor
  129. def apply_single(self, pred, tagmap):
  130. if tagmap.numpy()[:, :, 3].sum() == 0:
  131. return (paddle.zeros([1]), paddle.zeros([1]))
  132. nonzero = paddle.nonzero(tagmap[:, :, 3] > 0)
  133. if nonzero.shape[0] == 0:
  134. return (paddle.zeros([1]), paddle.zeros([1]))
  135. p_inds = paddle.unique(nonzero[:, 0])
  136. num_person = p_inds.shape[0]
  137. if num_person == 0:
  138. return (paddle.zeros([1]), paddle.zeros([1]))
  139. pull = 0
  140. tagpull_num = 0
  141. embs_all = []
  142. person_unvalid = 0
  143. for person_idx in p_inds.numpy():
  144. valid_single = tagmap[person_idx.item()]
  145. validkpts = paddle.nonzero(valid_single[:, 3] > 0)
  146. valid_single = paddle.index_select(valid_single, validkpts)
  147. emb = paddle.gather_nd(pred, valid_single[:, :3])
  148. if emb.shape[0] == 1:
  149. person_unvalid += 1
  150. mean = paddle.mean(emb, axis=0)
  151. embs_all.append(mean)
  152. pull += paddle.mean(paddle.pow(emb - mean, 2), axis=0)
  153. tagpull_num += emb.shape[0]
  154. pull /= max(num_person - person_unvalid, 1)
  155. if num_person < 2:
  156. return pull, paddle.zeros([1])
  157. embs_all = paddle.stack(embs_all)
  158. A = embs_all.expand([num_person, num_person])
  159. B = A.transpose([1, 0])
  160. diff = A - B
  161. diff = paddle.pow(diff, 2)
  162. push = paddle.exp(-diff)
  163. push = paddle.sum(push) - num_person
  164. push /= 2 * num_person * (num_person - 1)
  165. return pull, push
  166. def __call__(self, preds, tagmaps):
  167. bs = preds.shape[0]
  168. losses = [
  169. self.apply_single(preds[i:i + 1].squeeze(),
  170. tagmaps[i:i + 1].squeeze()) for i in range(bs)
  171. ]
  172. pull = self.pull_factor * sum(loss[0] for loss in losses) / len(losses)
  173. push = self.push_factor * sum(loss[1] for loss in losses) / len(losses)
  174. return pull, push
  175. class ZipLoss(object):
  176. def __init__(self, loss_funcs):
  177. super(ZipLoss, self).__init__()
  178. self.loss_funcs = loss_funcs
  179. def __call__(self, inputs, targets):
  180. assert len(self.loss_funcs) == len(targets) >= len(inputs)
  181. def zip_repeat(*args):
  182. longest = max(map(len, args))
  183. filled = [islice(cycle(x), longest) for x in args]
  184. return zip(*filled)
  185. return tuple(
  186. fn(x, y)
  187. for x, y, fn in zip_repeat(inputs, targets, self.loss_funcs))
  188. def recursive_sum(inputs):
  189. if isinstance(inputs, abc.Sequence):
  190. return sum([recursive_sum(x) for x in inputs])
  191. return inputs