keypoint_utils.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302
  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. import cv2
  15. import numpy as np
  16. def get_affine_mat_kernel(h, w, s, inv=False):
  17. if w < h:
  18. w_ = s
  19. h_ = int(np.ceil((s / w * h) / 64.) * 64)
  20. scale_w = w
  21. scale_h = h_ / w_ * w
  22. else:
  23. h_ = s
  24. w_ = int(np.ceil((s / h * w) / 64.) * 64)
  25. scale_h = h
  26. scale_w = w_ / h_ * h
  27. center = np.array([np.round(w / 2.), np.round(h / 2.)])
  28. size_resized = (w_, h_)
  29. trans = get_affine_transform(
  30. center, np.array([scale_w, scale_h]), 0, size_resized, inv=inv)
  31. return trans, size_resized
  32. def get_affine_transform(center,
  33. input_size,
  34. rot,
  35. output_size,
  36. shift=(0., 0.),
  37. inv=False):
  38. """Get the affine transform matrix, given the center/scale/rot/output_size.
  39. Args:
  40. center (np.ndarray[2, ]): Center of the bounding box (x, y).
  41. scale (np.ndarray[2, ]): Scale of the bounding box
  42. wrt [width, height].
  43. rot (float): Rotation angle (degree).
  44. output_size (np.ndarray[2, ]): Size of the destination heatmaps.
  45. shift (0-100%): Shift translation ratio wrt the width/height.
  46. Default (0., 0.).
  47. inv (bool): Option to inverse the affine transform direction.
  48. (inv=False: src->dst or inv=True: dst->src)
  49. Returns:
  50. np.ndarray: The transform matrix.
  51. """
  52. assert len(center) == 2
  53. assert len(input_size) == 2
  54. assert len(output_size) == 2
  55. assert len(shift) == 2
  56. scale_tmp = input_size
  57. shift = np.array(shift)
  58. src_w = scale_tmp[0]
  59. dst_w = output_size[0]
  60. dst_h = output_size[1]
  61. rot_rad = np.pi * rot / 180
  62. src_dir = rotate_point([0., src_w * -0.5], rot_rad)
  63. dst_dir = np.array([0., dst_w * -0.5])
  64. src = np.zeros((3, 2), dtype=np.float32)
  65. src[0, :] = center + scale_tmp * shift
  66. src[1, :] = center + src_dir + scale_tmp * shift
  67. src[2, :] = _get_3rd_point(src[0, :], src[1, :])
  68. dst = np.zeros((3, 2), dtype=np.float32)
  69. dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
  70. dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
  71. dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
  72. if inv:
  73. trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
  74. else:
  75. trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
  76. return trans
  77. def _get_3rd_point(a, b):
  78. """To calculate the affine matrix, three pairs of points are required. This
  79. function is used to get the 3rd point, given 2D points a & b.
  80. The 3rd point is defined by rotating vector `a - b` by 90 degrees
  81. anticlockwise, using b as the rotation center.
  82. Args:
  83. a (np.ndarray): point(x,y)
  84. b (np.ndarray): point(x,y)
  85. Returns:
  86. np.ndarray: The 3rd point.
  87. """
  88. assert len(a) == 2
  89. assert len(b) == 2
  90. direction = a - b
  91. third_pt = b + np.array([-direction[1], direction[0]], dtype=np.float32)
  92. return third_pt
  93. def rotate_point(pt, angle_rad):
  94. """Rotate a point by an angle.
  95. Args:
  96. pt (list[float]): 2 dimensional point to be rotated
  97. angle_rad (float): rotation angle by radian
  98. Returns:
  99. list[float]: Rotated point.
  100. """
  101. assert len(pt) == 2
  102. sn, cs = np.sin(angle_rad), np.cos(angle_rad)
  103. new_x = pt[0] * cs - pt[1] * sn
  104. new_y = pt[0] * sn + pt[1] * cs
  105. rotated_pt = [new_x, new_y]
  106. return rotated_pt
  107. def transpred(kpts, h, w, s):
  108. trans, _ = get_affine_mat_kernel(h, w, s, inv=True)
  109. return warp_affine_joints(kpts[..., :2].copy(), trans)
  110. def warp_affine_joints(joints, mat):
  111. """Apply affine transformation defined by the transform matrix on the
  112. joints.
  113. Args:
  114. joints (np.ndarray[..., 2]): Origin coordinate of joints.
  115. mat (np.ndarray[3, 2]): The affine matrix.
  116. Returns:
  117. matrix (np.ndarray[..., 2]): Result coordinate of joints.
  118. """
  119. joints = np.array(joints)
  120. shape = joints.shape
  121. joints = joints.reshape(-1, 2)
  122. return np.dot(np.concatenate(
  123. (joints, joints[:, 0:1] * 0 + 1), axis=1),
  124. mat.T).reshape(shape)
  125. def affine_transform(pt, t):
  126. new_pt = np.array([pt[0], pt[1], 1.]).T
  127. new_pt = np.dot(t, new_pt)
  128. return new_pt[:2]
  129. def transform_preds(coords, center, scale, output_size):
  130. target_coords = np.zeros(coords.shape)
  131. trans = get_affine_transform(center, scale * 200, 0, output_size, inv=1)
  132. for p in range(coords.shape[0]):
  133. target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
  134. return target_coords
  135. def oks_iou(g, d, a_g, a_d, sigmas=None, in_vis_thre=None):
  136. if not isinstance(sigmas, np.ndarray):
  137. sigmas = np.array([
  138. .26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07,
  139. .87, .87, .89, .89
  140. ]) / 10.0
  141. vars = (sigmas * 2)**2
  142. xg = g[0::3]
  143. yg = g[1::3]
  144. vg = g[2::3]
  145. ious = np.zeros((d.shape[0]))
  146. for n_d in range(0, d.shape[0]):
  147. xd = d[n_d, 0::3]
  148. yd = d[n_d, 1::3]
  149. vd = d[n_d, 2::3]
  150. dx = xd - xg
  151. dy = yd - yg
  152. e = (dx**2 + dy**2) / vars / ((a_g + a_d[n_d]) / 2 + np.spacing(1)) / 2
  153. if in_vis_thre is not None:
  154. ind = list(vg > in_vis_thre) and list(vd > in_vis_thre)
  155. e = e[ind]
  156. ious[n_d] = np.sum(np.exp(-e)) / e.shape[0] if e.shape[0] != 0 else 0.0
  157. return ious
  158. def oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
  159. """greedily select boxes with high confidence and overlap with current maximum <= thresh
  160. rule out overlap >= thresh
  161. Args:
  162. kpts_db (list): The predicted keypoints within the image
  163. thresh (float): The threshold to select the boxes
  164. sigmas (np.array): The variance to calculate the oks iou
  165. Default: None
  166. in_vis_thre (float): The threshold to select the high confidence boxes
  167. Default: None
  168. Return:
  169. keep (list): indexes to keep
  170. """
  171. if len(kpts_db) == 0:
  172. return []
  173. scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
  174. kpts = np.array(
  175. [kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
  176. areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
  177. order = scores.argsort()[::-1]
  178. keep = []
  179. while order.size > 0:
  180. i = order[0]
  181. keep.append(i)
  182. oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]],
  183. sigmas, in_vis_thre)
  184. inds = np.where(oks_ovr <= thresh)[0]
  185. order = order[inds + 1]
  186. return keep
  187. def rescore(overlap, scores, thresh, type='gaussian'):
  188. assert overlap.shape[0] == scores.shape[0]
  189. if type == 'linear':
  190. inds = np.where(overlap >= thresh)[0]
  191. scores[inds] = scores[inds] * (1 - overlap[inds])
  192. else:
  193. scores = scores * np.exp(-overlap**2 / thresh)
  194. return scores
  195. def soft_oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
  196. """greedily select boxes with high confidence and overlap with current maximum <= thresh
  197. rule out overlap >= thresh
  198. Args:
  199. kpts_db (list): The predicted keypoints within the image
  200. thresh (float): The threshold to select the boxes
  201. sigmas (np.array): The variance to calculate the oks iou
  202. Default: None
  203. in_vis_thre (float): The threshold to select the high confidence boxes
  204. Default: None
  205. Return:
  206. keep (list): indexes to keep
  207. """
  208. if len(kpts_db) == 0:
  209. return []
  210. scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
  211. kpts = np.array(
  212. [kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
  213. areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
  214. order = scores.argsort()[::-1]
  215. scores = scores[order]
  216. # max_dets = order.size
  217. max_dets = 20
  218. keep = np.zeros(max_dets, dtype=np.intp)
  219. keep_cnt = 0
  220. while order.size > 0 and keep_cnt < max_dets:
  221. i = order[0]
  222. oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]],
  223. sigmas, in_vis_thre)
  224. order = order[1:]
  225. scores = rescore(oks_ovr, scores[1:], thresh)
  226. tmp = scores.argsort()[::-1]
  227. order = order[tmp]
  228. scores = scores[tmp]
  229. keep[keep_cnt] = i
  230. keep_cnt += 1
  231. keep = keep[:keep_cnt]
  232. return keep