基于EfficientHRNet的Associative Embedding姿态关键点分组技术咨询

Associative Embeddings的核心逻辑很直接：每个检测到的关键点会对应输出一个低维embedding向量，同一人体的关键点向量相似度极高，不同人体的向量差异明显。我们只需要基于这些向量的相似度做聚类，就能把零散的关键点归整为单个人的完整姿态。

以下是针对已提取关键点的分组代码实现，假设你已经从EfficientHRNet的输出中拿到了每个关键点的坐标、置信度和对应的embedding向量：

import numpy as np
from scipy.spatial.distance import cdist
from sklearn.cluster import DBSCAN

def group_keypoints_by_associative_embedding(keypoints, embedding_threshold=0.5, eps=0.3, min_samples=3):
    """
    Args:
        keypoints: 列表，每个元素为 [x, y, score, embedding_vector]，embedding_vector为低维数组（如16维）
        embedding_threshold: 关键点置信度过滤阈值，低于该值的点直接丢弃
        eps: DBSCAN聚类的邻域半径（基于余弦相似度转换后的距离）
        min_samples: DBSCAN聚类的最小样本数，即一个人体至少需要多少个关键点才会被识别
    
    Returns:
        grouped_keypoints: 列表，每个元素为同一人体的所有关键点列表
    """
    # 过滤低置信度的关键点
    filtered_kps = [kp for kp in keypoints if kp[2] > embedding_threshold]
    if not filtered_kps:
        return []
    
    # 提取所有embedding向量
    embeddings = np.array([kp[3] for kp in filtered_kps])
    # 计算余弦相似度矩阵，转换为距离形式（1 - 相似度）
    cos_dist = 1 - cdist(embeddings, embeddings, metric='cosine')
    
    # 用DBSCAN聚类，基于预计算的余弦距离分组
    dbscan = DBSCAN(eps=eps, min_samples=min_samples, metric='precomputed')
    cluster_labels = dbscan.fit_predict(cos_dist)
    
    # 根据聚类标签分组关键点，忽略噪声点（标签为-1）
    grouped_keypoints = []
    unique_labels = np.unique(cluster_labels)
    for label in unique_labels:
        if label == -1:
            continue
        group = [filtered_kps[i] for i, lbl in enumerate(cluster_labels) if lbl == label]
        grouped_keypoints.append(group)
    
    return grouped_keypoints

# ---------------------- 示例用法 ----------------------
# 模拟从Heatmap提取的关键点格式
sample_keypoints = [
    [100, 200, 0.95, np.random.rand(16)],  # 人体1的关键点1
    [120, 250, 0.92, np.random.rand(16)],  # 人体1的关键点2
    [300, 180, 0.90, np.random.rand(16)],  # 人体2的关键点1
    [320, 230, 0.88, np.random.rand(16)],  # 人体2的关键点2
    [50, 100, 0.40, np.random.rand(16)],   # 低置信度噪声点
]

# 执行分组
groups = group_keypoints_by_associative_embedding(sample_keypoints)
for idx, group in enumerate(groups):
    print(f"第{idx+1}个人体的关键点：")
    for kp in group:
        print(f"坐标：({kp[0]}, {kp[1]}), 置信度：{kp[2]:.2f}")

关键逻辑说明 #

1. 置信度过滤：先丢弃置信度低的关键点，避免噪声干扰聚类效果。
2. 相似度计算：用余弦相似度衡量embedding向量的相似性，转换为距离后适配DBSCAN的输入要求。
3. DBSCAN聚类：无需预设人体数量，自动根据embedding的密度分组，适合多人姿态场景；eps和min_samples需根据数据集调整，通常eps在0.2-0.5之间，min_samples设为人体关键点总数的1/3左右（比如COCO的17个关键点，可设为5-6）。

如果不想依赖sklearn，也可以实现简单的阈值匹配逻辑（遍历每个点，将相似度高于阈值的点归为一组），但DBSCAN的鲁棒性更强，更适合复杂场景。

内容的提问来源于stack exchange，提问作者Ajesh Mon R