Source code for mmpose.codecs.associative_embedding
# Copyright (c) OpenMMLab. All rights reserved.
from itertools import product
from typing import Any, List, Optional, Tuple
import numpy as np
import torch
from munkres import Munkres
from torch import Tensor
from mmpose.registry import KEYPOINT_CODECS
from mmpose.utils.tensor_utils import to_numpy
from .base import BaseKeypointCodec
from .utils import (batch_heatmap_nms, generate_gaussian_heatmaps,
generate_udp_gaussian_heatmaps, refine_keypoints,
refine_keypoints_dark_udp)
def _py_max_match(scores):
"""Apply munkres algorithm to get the best match.
Args:
scores(np.ndarray): cost matrix.
Returns:
np.ndarray: best match.
"""
m = Munkres()
tmp = m.compute(scores)
tmp = np.array(tmp).astype(int)
return tmp
def _group_keypoints_by_tags(vals: np.ndarray,
tags: np.ndarray,
locs: np.ndarray,
keypoint_order: List[int],
val_thr: float,
tag_thr: float = 1.0,
max_groups: Optional[int] = None) -> np.ndarray:
"""Group the keypoints by tags using Munkres algorithm.
Note:
- keypoint number: K
- candidate number: M
- tag dimenssion: L
- coordinate dimension: D
- group number: G
Args:
vals (np.ndarray): The heatmap response values of keypoints in shape
(K, M)
tags (np.ndarray): The tags of the keypoint candidates in shape
(K, M, L)
locs (np.ndarray): The locations of the keypoint candidates in shape
(K, M, D)
keypoint_order (List[int]): The grouping order of the keypoints.
The groupping usually starts from a keypoints around the head and
torso, and gruadually moves out to the limbs
val_thr (float): The threshold of the keypoint response value
tag_thr (float): The maximum allowed tag distance when matching a
keypoint to a group. A keypoint with larger tag distance to any
of the existing groups will initializes a new group
max_groups (int, optional): The maximum group number. ``None`` means
no limitation. Defaults to ``None``
Returns:
np.ndarray: grouped keypoints in shape (G, K, D+1), where the last
dimenssion is the concatenated keypoint coordinates and scores.
"""
tag_k, loc_k, val_k = tags, locs, vals
K, M, D = locs.shape
assert vals.shape == tags.shape[:2] == (K, M)
assert len(keypoint_order) == K
default_ = np.zeros((K, 3 + tag_k.shape[2]), dtype=np.float32)
joint_dict = {}
tag_dict = {}
for i in range(K):
idx = keypoint_order[i]
tags = tag_k[idx]
joints = np.concatenate((loc_k[idx], val_k[idx, :, None], tags), 1)
mask = joints[:, 2] > val_thr
tags = tags[mask] # shape: [M, L]
joints = joints[mask] # shape: [M, 3 + L], 3: x, y, val
if joints.shape[0] == 0:
continue
if i == 0 or len(joint_dict) == 0:
for tag, joint in zip(tags, joints):
key = tag[0]
joint_dict.setdefault(key, np.copy(default_))[idx] = joint
tag_dict[key] = [tag]
else:
# shape: [M]
grouped_keys = list(joint_dict.keys())
# shape: [M, L]
grouped_tags = [np.mean(tag_dict[i], axis=0) for i in grouped_keys]
# shape: [M, M, L]
diff = joints[:, None, 3:] - np.array(grouped_tags)[None, :, :]
# shape: [M, M]
diff_normed = np.linalg.norm(diff, ord=2, axis=2)
diff_saved = np.copy(diff_normed)
diff_normed = np.round(diff_normed) * 100 - joints[:, 2:3]
num_added = diff.shape[0]
num_grouped = diff.shape[1]
if num_added > num_grouped:
diff_normed = np.concatenate(
(diff_normed,
np.zeros((num_added, num_added - num_grouped),
dtype=np.float32) + 1e10),
axis=1)
pairs = _py_max_match(diff_normed)
for row, col in pairs:
if (row < num_added and col < num_grouped
and diff_saved[row][col] < tag_thr):
key = grouped_keys[col]
joint_dict[key][idx] = joints[row]
tag_dict[key].append(tags[row])
else:
key = tags[row][0]
joint_dict.setdefault(key, np.copy(default_))[idx] = \
joints[row]
tag_dict[key] = [tags[row]]
joint_dict_keys = list(joint_dict.keys())[:max_groups]
if joint_dict_keys:
results = np.array([joint_dict[i]
for i in joint_dict_keys]).astype(np.float32)
results = results[..., :D + 1]
else:
results = np.empty((0, K, D + 1), dtype=np.float32)
return results
[docs]@KEYPOINT_CODECS.register_module()
class AssociativeEmbedding(BaseKeypointCodec):
"""Encode/decode keypoints with the method introduced in "Associative
Embedding". This is an asymmetric codec, where the keypoints are
represented as gaussian heatmaps and position indices during encoding, and
restored from predicted heatmaps and group tags.
See the paper `Associative Embedding: End-to-End Learning for Joint
Detection and Grouping`_ by Newell et al (2017) for details
Note:
- instance number: N
- keypoint number: K
- keypoint dimension: D
- embedding tag dimension: L
- image size: [w, h]
- heatmap size: [W, H]
Encoded:
- heatmaps (np.ndarray): The generated heatmap in shape (K, H, W)
where [W, H] is the `heatmap_size`
- keypoint_indices (np.ndarray): The keypoint position indices in shape
(N, K, 2). Each keypoint's index is [i, v], where i is the position
index in the heatmap (:math:`i=y*w+x`) and v is the visibility
- keypoint_weights (np.ndarray): The target weights in shape (N, K)
Args:
input_size (tuple): Image size in [w, h]
heatmap_size (tuple): Heatmap size in [W, H]
sigma (float): The sigma value of the Gaussian heatmap
use_udp (bool): Whether use unbiased data processing. See
`UDP (CVPR 2020)`_ for details. Defaults to ``False``
decode_keypoint_order (List[int]): The grouping order of the
keypoint indices. The groupping usually starts from a keypoints
around the head and torso, and gruadually moves out to the limbs
decode_keypoint_thr (float): The threshold of keypoint response value
in heatmaps. Defaults to 0.1
decode_tag_thr (float): The maximum allowed tag distance when matching
a keypoint to a group. A keypoint with larger tag distance to any
of the existing groups will initializes a new group. Defaults to
1.0
decode_nms_kernel (int): The kernel size of the NMS during decoding,
which should be an odd integer. Defaults to 5
decode_gaussian_kernel (int): The kernel size of the Gaussian blur
during decoding, which should be an odd integer. It is only used
when ``self.use_udp==True``. Defaults to 3
decode_topk (int): The number top-k candidates of each keypoints that
will be retrieved from the heatmaps during dedocding. Defaults to
20
decode_max_instances (int, optional): The maximum number of instances
to decode. ``None`` means no limitation to the instance number.
Defaults to ``None``
.. _`Associative Embedding: End-to-End Learning for Joint Detection and
Grouping`: https://arxiv.org/abs/1611.05424
.. _`UDP (CVPR 2020)`: https://arxiv.org/abs/1911.07524
"""
def __init__(
self,
input_size: Tuple[int, int],
heatmap_size: Tuple[int, int],
sigma: Optional[float] = None,
use_udp: bool = False,
decode_keypoint_order: List[int] = [],
decode_nms_kernel: int = 5,
decode_gaussian_kernel: int = 3,
decode_keypoint_thr: float = 0.1,
decode_tag_thr: float = 1.0,
decode_topk: int = 30,
decode_center_shift=0.0,
decode_max_instances: Optional[int] = None,
) -> None:
super().__init__()
self.input_size = input_size
self.heatmap_size = heatmap_size
self.use_udp = use_udp
self.decode_nms_kernel = decode_nms_kernel
self.decode_gaussian_kernel = decode_gaussian_kernel
self.decode_keypoint_thr = decode_keypoint_thr
self.decode_tag_thr = decode_tag_thr
self.decode_topk = decode_topk
self.decode_center_shift = decode_center_shift
self.decode_max_instances = decode_max_instances
self.decode_keypoint_order = decode_keypoint_order.copy()
if self.use_udp:
self.scale_factor = ((np.array(input_size) - 1) /
(np.array(heatmap_size) - 1)).astype(
np.float32)
else:
self.scale_factor = (np.array(input_size) /
heatmap_size).astype(np.float32)
if sigma is None:
sigma = (heatmap_size[0] * heatmap_size[1])**0.5 / 64
self.sigma = sigma
[docs] def encode(
self,
keypoints: np.ndarray,
keypoints_visible: Optional[np.ndarray] = None
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Encode keypoints into heatmaps and position indices. Note that the
original keypoint coordinates should be in the input image space.
Args:
keypoints (np.ndarray): Keypoint coordinates in shape (N, K, D)
keypoints_visible (np.ndarray): Keypoint visibilities in shape
(N, K)
Returns:
dict:
- heatmaps (np.ndarray): The generated heatmap in shape
(K, H, W) where [W, H] is the `heatmap_size`
- keypoint_indices (np.ndarray): The keypoint position indices
in shape (N, K, 2). Each keypoint's index is [i, v], where i
is the position index in the heatmap (:math:`i=y*w+x`) and v
is the visibility
- keypoint_weights (np.ndarray): The target weights in shape
(N, K)
"""
if keypoints_visible is None:
keypoints_visible = np.ones(keypoints.shape[:2], dtype=np.float32)
# keypoint coordinates in heatmap
_keypoints = keypoints / self.scale_factor
if self.use_udp:
heatmaps, keypoint_weights = generate_udp_gaussian_heatmaps(
heatmap_size=self.heatmap_size,
keypoints=_keypoints,
keypoints_visible=keypoints_visible,
sigma=self.sigma)
else:
heatmaps, keypoint_weights = generate_gaussian_heatmaps(
heatmap_size=self.heatmap_size,
keypoints=_keypoints,
keypoints_visible=keypoints_visible,
sigma=self.sigma)
keypoint_indices = self._encode_keypoint_indices(
heatmap_size=self.heatmap_size,
keypoints=_keypoints,
keypoints_visible=keypoints_visible)
encoded = dict(
heatmaps=heatmaps,
keypoint_indices=keypoint_indices,
keypoint_weights=keypoint_weights)
return encoded
def _encode_keypoint_indices(self, heatmap_size: Tuple[int, int],
keypoints: np.ndarray,
keypoints_visible: np.ndarray) -> np.ndarray:
w, h = heatmap_size
N, K, _ = keypoints.shape
keypoint_indices = np.zeros((N, K, 2), dtype=np.int64)
for n, k in product(range(N), range(K)):
x, y = (keypoints[n, k] + 0.5).astype(np.int64)
index = y * w + x
vis = (keypoints_visible[n, k] > 0.5 and 0 <= x < w and 0 <= y < h)
keypoint_indices[n, k] = [index, vis]
return keypoint_indices
def _get_batch_topk(self, batch_heatmaps: Tensor, batch_tags: Tensor,
k: int):
"""Get top-k response values from the heatmaps and corresponding tag
values from the tagging heatmaps.
Args:
batch_heatmaps (Tensor): Keypoint detection heatmaps in shape
(B, K, H, W)
batch_tags (Tensor): Tagging heatmaps in shape (B, C, H, W), where
the tag dim C is 2*K when using flip testing, or K otherwise
k (int): The number of top responses to get
Returns:
tuple:
- topk_vals (Tensor): Top-k response values of each heatmap in
shape (B, K, Topk)
- topk_tags (Tensor): The corresponding embedding tags of the
top-k responses, in shape (B, K, Topk, L)
- topk_locs (Tensor): The location of the top-k responses in each
heatmap, in shape (B, K, Topk, 2) where last dimension
represents x and y coordinates
"""
B, K, H, W = batch_heatmaps.shape
L = batch_tags.shape[1] // K
# shape of topk_val, top_indices: (B, K, TopK)
topk_vals, topk_indices = batch_heatmaps.flatten(-2, -1).topk(
k, dim=-1)
topk_tags_per_kpts = [
torch.gather(_tag, dim=2, index=topk_indices)
for _tag in torch.unbind(batch_tags.view(B, L, K, H * W), dim=1)
]
topk_tags = torch.stack(topk_tags_per_kpts, dim=-1) # (B, K, TopK, L)
topk_locs = torch.stack([topk_indices % W, topk_indices // W],
dim=-1) # (B, K, TopK, 2)
return topk_vals, topk_tags, topk_locs
def _group_keypoints(self, batch_vals: np.ndarray, batch_tags: np.ndarray,
batch_locs: np.ndarray):
"""Group keypoints into groups (each represents an instance) by tags.
Args:
batch_vals (Tensor): Heatmap response values of keypoint
candidates in shape (B, K, Topk)
batch_tags (Tensor): Tags of keypoint candidates in shape
(B, K, Topk, L)
batch_locs (Tensor): Locations of keypoint candidates in shape
(B, K, Topk, 2)
Returns:
List[np.ndarray]: Grouping results of a batch, each element is a
np.ndarray (in shape [N, K, D+1]) that contains the groups
detected in an image, including both keypoint coordinates and
scores.
"""
def _group_func(inputs: Tuple):
vals, tags, locs = inputs
return _group_keypoints_by_tags(
vals,
tags,
locs,
keypoint_order=self.decode_keypoint_order,
val_thr=self.decode_keypoint_thr,
tag_thr=self.decode_tag_thr,
max_groups=self.decode_max_instances)
_results = map(_group_func, zip(batch_vals, batch_tags, batch_locs))
results = list(_results)
return results
def _fill_missing_keypoints(self, keypoints: np.ndarray,
keypoint_scores: np.ndarray,
heatmaps: np.ndarray, tags: np.ndarray):
"""Fill the missing keypoints in the initial predictions.
Args:
keypoints (np.ndarray): Keypoint predictions in shape (N, K, D)
keypoint_scores (np.ndarray): Keypint score predictions in shape
(N, K), in which 0 means the corresponding keypoint is
missing in the initial prediction
heatmaps (np.ndarry): Heatmaps in shape (K, H, W)
tags (np.ndarray): Tagging heatmaps in shape (C, H, W) where
C=L*K
Returns:
tuple:
- keypoints (np.ndarray): Keypoint predictions with missing
ones filled
- keypoint_scores (np.ndarray): Keypoint score predictions with
missing ones filled
"""
N, K = keypoints.shape[:2]
H, W = heatmaps.shape[1:]
L = tags.shape[0] // K
keypoint_tags = [tags[k::K] for k in range(K)]
for n in range(N):
# Calculate the instance tag (mean tag of detected keypoints)
_tag = []
for k in range(K):
if keypoint_scores[n, k] > 0:
x, y = keypoints[n, k, :2].astype(np.int64)
x = np.clip(x, 0, W - 1)
y = np.clip(y, 0, H - 1)
_tag.append(keypoint_tags[k][:, y, x])
tag = np.mean(_tag, axis=0)
tag = tag.reshape(L, 1, 1)
# Search maximum response of the missing keypoints
for k in range(K):
if keypoint_scores[n, k] > 0:
continue
dist_map = np.linalg.norm(
keypoint_tags[k] - tag, ord=2, axis=0)
cost_map = np.round(dist_map) * 100 - heatmaps[k] # H, W
y, x = np.unravel_index(np.argmin(cost_map), shape=(H, W))
keypoints[n, k] = [x, y]
keypoint_scores[n, k] = heatmaps[k, y, x]
return keypoints, keypoint_scores
[docs] def batch_decode(self, batch_heatmaps: Tensor, batch_tags: Tensor
) -> Tuple[List[np.ndarray], List[np.ndarray]]:
"""Decode the keypoint coordinates from a batch of heatmaps and tagging
heatmaps. The decoded keypoint coordinates are in the input image
space.
Args:
batch_heatmaps (Tensor): Keypoint detection heatmaps in shape
(B, K, H, W)
batch_tags (Tensor): Tagging heatmaps in shape (B, C, H, W), where
:math:`C=L*K`
Returns:
tuple:
- batch_keypoints (List[np.ndarray]): Decoded keypoint coordinates
of the batch, each is in shape (N, K, D)
- batch_scores (List[np.ndarray]): Decoded keypoint scores of the
batch, each is in shape (N, K). It usually represents the
confidience of the keypoint prediction
"""
B, _, H, W = batch_heatmaps.shape
assert batch_tags.shape[0] == B and batch_tags.shape[2:4] == (H, W), (
f'Mismatched shapes of heatmap ({batch_heatmaps.shape}) and '
f'tagging map ({batch_tags.shape})')
# Heatmap NMS
batch_heatmaps_peak = batch_heatmap_nms(batch_heatmaps,
self.decode_nms_kernel)
# Get top-k in each heatmap and and convert to numpy
batch_topk_vals, batch_topk_tags, batch_topk_locs = to_numpy(
self._get_batch_topk(
batch_heatmaps_peak, batch_tags, k=self.decode_topk))
# Group keypoint candidates into groups (instances)
batch_groups = self._group_keypoints(batch_topk_vals, batch_topk_tags,
batch_topk_locs)
# Convert to numpy
batch_heatmaps_np = to_numpy(batch_heatmaps)
batch_tags_np = to_numpy(batch_tags)
# Refine the keypoint prediction
batch_keypoints = []
batch_keypoint_scores = []
batch_instance_scores = []
for i, (groups, heatmaps, tags) in enumerate(
zip(batch_groups, batch_heatmaps_np, batch_tags_np)):
keypoints, scores = groups[..., :-1], groups[..., -1]
instance_scores = scores.mean(axis=-1)
if keypoints.size > 0:
# refine keypoint coordinates according to heatmap distribution
if self.use_udp:
keypoints = refine_keypoints_dark_udp(
keypoints,
heatmaps,
blur_kernel_size=self.decode_gaussian_kernel)
else:
keypoints = refine_keypoints(keypoints, heatmaps)
keypoints += self.decode_center_shift * \
(scores > 0).astype(keypoints.dtype)[..., None]
# identify missing keypoints
keypoints, scores = self._fill_missing_keypoints(
keypoints, scores, heatmaps, tags)
batch_keypoints.append(keypoints)
batch_keypoint_scores.append(scores)
batch_instance_scores.append(instance_scores)
# restore keypoint scale
batch_keypoints = [
kpts * self.scale_factor for kpts in batch_keypoints
]
return batch_keypoints, batch_keypoint_scores, batch_instance_scores