AlphaPose Pytorch 代码详解(一):predict

前言

代码地址:AlphaPose-Pytorch版

本文以图像 1.jpg(854x480)为例对整个预测过程的各个细节进行解读并记录

python demo.py --indir examples/demo --outdir examples/res --save_img

在这里插入图片描述

1. YOLO

1.1 图像预处理

  • cv2读取BGR图像 img [480,854,3] (h,w,c)
  • 按参数 inp_dim=608,将图像保持长宽比缩放至 [341,608,3],并以数值为128做padding至 [608,608,3]
  • BGR转RGB、维度处理、转tensor、数据类型转float、数值除以255做归一化,[1,3,608,608] (b,c,h,w)

输出:

img:		[1,3,608,608]	(yolo输入图像)
orig_img:	[480,854,3]		(原始BGR图像)
im_name:	'examples/demo/1.jpg'
im_dim_list:[854,480,854,480](原图尺寸, 用于把yolo的输出坐标转换到原图坐标系)

1.2 yolo 模型推理

输入:img[1,3,608,608]
输出:pred[1,22743,85]

1 = b a t c h s i z e mathrm{1=batchsize} 1=batchsize

22743 = [ ( 608 / 32 ) 2 + ( 608 / 16 ) 2 + ( 608 / 8 ) 2 ] × 3 22743=[(608/32)^2+(608/16)^2+(608/8)^2]times3 22743=[(608/32)2+(608/16)2+(608/8)2]×3

85 = [ x , y , w , h , c o n f , 80 c l a s s e s ] mathrm{85=[x,y,w,h,conf,80classes]} 85=[x,y,w,h,conf,80classes]

1.3 输出后处理

(1)第一阶段

  • 坐标 xywhxyxy
  • batchsize 循环,[22743,85][22743,7] 7 = [ x , y , x , y , c o n f , c l a s s   s c o r e , c l a s s ] mathrm{7=[x,y,x,y,conf,class score,class]} 7=[x,y,x,y,conf,class score,class],即80类得分转换为得分最高的类别索引及其得分
  • 去除 c o n f ≤ 0.05 mathrm{confle0.05} conf0.05 的项 [37,7]
  • 保留类别为人的项,并且按 conf 从高到低排序,得到结果 img_pred[19,7]
  • nms去除重复目标 [19,7]->[6,7]
  • 添加 batch_idx 这里batchsize为1所以都是0,[6,7]->[6,8] 8 = [ b a t c h   i d x , x , y , x , y , c o n f , c l a s s   s c o r e , c l a s s ] mathrm{8=[batch idx, x,y,x,y,conf,class score,class]} 8=[batch idx,x,y,x,y,conf,class score,class]
  • 坐标数值转换,从 [608,608] 转到原图坐标 [854,480],并把坐标 clamp[0,w] [0,h] 之间

输出:

orig_img:[480,854,3]	(原始BGR图像)
im_name:'examples/demo/1.jpg'
boxes:	[6,4]	(x,y,x,y)(原图坐标系)
scores:	[6,1]	(conf)

NMS 细节

  • nms_conf=0.6
  • img_pred 中的第一项放到结果中
  • 剩余所有项与结果中的最后一项计算iou,保留 iou<nms_conf 的项作为新的 img_pred
  • 循环直到 img_pred 中没有目标
  • 当经过nms后的目标数量大于100个,会把 nms_conf-0.05,从最初的 img_pred 开始重新进行nms

(2)第二阶段

  • 原始图像 orig_img [480,854,3] BGR转RGB、维度处理、转tensor、数据类型转float、数值除以255做归一化,得到 inp [3,480,854]
  • 对三通道做处理 inp[0].add_(-0.406), inp[1].add_(-0.457), inp[2].add_(-0.480)
  • 扩大 boxes 中目标框的范围,并把左上角坐标存入 pt1,右下角坐标存入 pt2
  • 根据 boxes 把每个目标从图中抠出来,通过保比例缩放+zero padding,统一成 [3,320,256] 大小的图像存入 inps

输出:

inps:	[6,3,320,256]	(检测目标的子图像,作为Alphapose的输入)
orig_img:[480,854,3]	(原始BGR图像)
im_name:'examples/demo/1.jpg'
boxes:	[6,4]	(x1,y1,x2,y2)(yolo原始输出,原图坐标系)
scores:	[6,1]	(yolo输出conf)
pt1:	[6,2]	(x1,y1)(yolo输出扩大后坐标,原图坐标系)
pt2:	[6,2]	(x2,y2)(yolo输出扩大后坐标,原图坐标系)

2. POSE

2.1 pose 模型推理

输入:inps[6,3,320,256]
输出:hm[6,17,80,64],即6个目标,每个目标17个关键点对应的热力图

2.2 输出后处理

(1)第一阶段:热力图转坐标

  • 获取 hm[6,17,80,64] 中每个关键点的热力图中最大值的索引 preds[6,17,2]
  • 由于 opt.matching=False,此处使用简单的后处理,源码如下
    preds 中某个索引 [x,y] 为例,取出其热力图中相邻的上下左右四个位置的值,并且分别在 x x x y y y 轴上往较高的方向偏移 0.25 0.25 0.25
    x x x 轴为例: p l e f t = h m [ y ] [ x − 1 ] p_mathrm{left}=mathrm{hm}[y][x-1] pleft=hm[y][x1] p r i g h t = h m [ y ] [ x + 1 ] p_mathrm{right}=mathrm{hm}[y][x+1] pright=hm[y][x+1],若 p l e f t > p r i g h t p_mathrm{left}>p_mathrm{right} pleft>pright x + 0.25 x+0.25 x+0.25,若 p l e f t = p r i g h t p_mathrm{left}=p_mathrm{right} pleft=pright x x x 保持不变
    最后会在所有的坐标值上 + 0.2 +0.2 +0.2
for i in range(preds.size(0)):
    for j in range(preds.size(1)):
        hm = hms[i][j]
        pX, pY = int(round(float(preds[i][j][0]))), int(round(float(preds[i][j][1])))
        if 0 < pX < opt.outputResW - 1 and 0 < pY < opt.outputResH - 1:
            diff = torch.Tensor((hm[pY][pX + 1] - hm[pY][pX - 1], hm[pY + 1][pX] - hm[pY - 1][pX]))
            preds[i][j] += diff.sign() * 0.25
preds += 0.2
  • 目前得到的坐标 preds[6,17,2] 是相对于输出分辨率 [80,64] 坐标系下的,转换到原图分辨率 [480,854] 坐标系下,得到 preds_tf[6,17,2]

输出:

preds:		[6,17,2]	(经过第二步偏移处理后的坐标,相对于热力图坐标系)
preds_tf:	[6,17,2]	(最终坐标,相对于原图坐标系)
maxval:		[6,17,1]	(热力图最大值)

(2)第二阶段:pose nms

输入:

ori_bboxs:		[6,4]		(yolo原始输出,原图坐标系)
ori_bbox_scores:[6,1]		(yolo输出conf)
ori_pose_preds:	[6,17,2]	(对应preds_tf,关键点坐标,原图坐标系)
ori_pose_scores:[6,17,1]	(对应maxval,热力图最大值)
  • 根据 bboxs 计算每个目标框的 w,h,选择每个目标框中的最大值 max(w,h) 并乘上 alpha=0.1 构成 ref_dists[6]
  • 根据 pose_scores 计算每个目标17个关键点得分的均值,得到 human_scores[6]
  • 开始循环,直到 human_scores 无目标
    1. 选择 human_scores 最高的目标,坐标和得分分别记为 pick_preds[17,2], pick_scores[17,1]
      全部的坐标和得分记为 all_preds[6,17,2], all_scores[6,17,1](此处命名方式与源码略有不同,以便于区分)
    2. 计算距离:final_dist[6] 目标的同类别关键点的距离,距离越近数值越大
      score_dists 计算位置距离非常近的同类别关键点的得分距离
      point_dist = e − d / 2.65 =e^{-d/2.65} =ed/2.65,因为 d ≥ 0 dge0 d0,所以 0 < p o i n t _ d i s t ≤ 1 0<mathrm{point_dist}le1 0<point_dist1 d d d 越小, p o i n t _ d i s t mathrm{point_dist} point_dist 越大,目标本身则最大全为1
    3. 计算关键点匹配数量:num_match_keypoints[6] 目标之间同类别关键点中距离较近的数量
    4. 去除多余目标:目标之间的距离超过阈值 or 目标之间距离相近的关键点数量超过阈值 → 判定为多余的目标。
      由于选出的目标本身也在其中,因此目标自身必然在去除的队伍中,如果除了自身还有目标被去除,那么会把额外的目标与自身的索引放在一起得到 merge_ids,这些目标相互之间距离很近,用于后续融合目标。
对应第2def get_parametric_distance(i, all_preds, all_scores):
	pick_preds, pick_scores = all_preds[i], all_scores[i]
	'计算坐标位置的欧氏距离 dist[6,17](同类别关键点之间的距离)'
	dist = torch.sqrt(torch.sum(torch.pow(pick_preds[np.newaxis, :] - all_preds, 2), dim=2))
	
	'计算dist<=1的点之间的得分距离 score_dists[6,17]'
	mask = (dist <= 1)
	score_dists = torch.zeros(all_preds.shape[0], 17)
	score_dists[mask] = torch.tanh(pick_scores[mask]/delta1) * torch.tanh(all_scores[mask]/delta1)  'delta1=1'
	
	'final_dist[6]'
	point_dist = torch.exp((-1) * dist / delta2)  'delta2=2.65'
	final_dist = torch.sum(score_dists, dim=1) + mu * torch.sum(point_dist, dim=1)  'mu=1.7'
	return final_dist
对应第3def PCK_match(pick_pred, all_preds, ref_dist):
	dist = torch.sqrt(torch.sum(torch.pow(pick_preds[np.newaxis, :] - all_preds, 2), dim=2))
	ref_dist = min(ref_dist, 7)
	num_match_keypoints = torch.sum(dist / ref_dist <= 1, dim=1)
	return num_match_keypoints
对应第4'gamma=22.48, matchThreds=5'
delete_ids = torch.from_numpy(np.arange(human_scores.shape[0]))[(final_dist > gamma) | (num_match_keypoints >= matchThreds)]

输出:

merge_ids:			[6,x]
preds_pick:			[6,17,2]
scores_pick:		[6,17,1]
bbox_scores_pick:	[6,1]
'''
这里的输出是从各个输入 orig_xxxx, 例如 ori_bbox_scores 中挑选出来的(nms后的目标)
只是在这个例子中, nms判断并没有重复的目标, 因此和原始输入保持一致

merge_ids 是一个列表, x代表每一项的长度, 本例中x都=1
如果nms判断存在重复目标, 那么会把这些目标在原始输入中的索引记录在 merge_ids 中, 此时x>1
在第三阶段中会把这些目标进行融合
'''

(3)第三阶段:融合与过滤

  • 去除17个关键点中最高得分 m a x _ s c o r e < s c o r e T h r e d s = 0.3 mathrm{max_score < scoreThreds = 0.3} max_score<scoreThreds=0.3 的目标
  • 融合目标,具体看下面代码,简单来说就是把距离比较近的关键点根据得分的高低作为权重,把坐标位置和得分进行加权求和作为融合后的目标
  • 去除融合后,17个关键点中最高得分 m a x _ s c o r e < s c o r e T h r e d s = 0.3 mathrm{max_score < scoreThreds = 0.3} max_score<scoreThreds=0.3 的目标
  • 根据能包含目标所有关键点的矩形框面积来过滤目标,1.5**2 * (xmax-xmin) * (ymax-ymin) < areaThres=0,具体为外接矩形长宽都乘1.5后计算面积,由于这里阈值为0,过滤基本无效
  • 最后会把所有关键点坐标数值 − 0.3 -0.3 0.3,并且根据关键点得分和目标框得分生成 proposal_score,具体见下面代码

此阶段过滤掉了一个目标,最终得到5个目标。

merge_pose, merge_score = p_merge_fast(preds_pick[j], ori_pose_preds[merge_id], ori_pose_scores[merge_id], ref_dists[pick[j]])

def p_merge_fast(ref_pose, cluster_preds, cluster_scores, ref_dist):
    '''
    Score-weighted pose merging
    INPUT:															本博客中文别称
        ref_pose:       reference pose          -- [17, 2]			挑选目标关键点
        cluster_preds:  redundant poses         -- [n, 17, 2]		多余目标关键点(挑选目标本身包含在多余目标中)
        cluster_scores: redundant poses score   -- [n, 17, 1]
        ref_dist:       reference scale         -- Constant
    OUTPUT:
        final_pose:     merged pose             -- [17, 2]
        final_score:    merged score            -- [17]
    '''
    
    '计算与多余目标关键点距离 dist[n,17]'
    dist = torch.sqrt(torch.sum(
        torch.pow(ref_pose[np.newaxis, :] - cluster_preds, 2),
        dim=2
    ))

    kp_num = 17
    '回顾一下, ref_dist是挑选目标的目标框的 max(h,w)*0.1'
    ref_dist = min(ref_dist, 15)

    mask = (dist <= ref_dist)
    final_pose = torch.zeros(kp_num, 2)
    final_score = torch.zeros(kp_num)

    if cluster_preds.dim() == 2:
        cluster_preds.unsqueeze_(0)
        cluster_scores.unsqueeze_(0)
    if mask.dim() == 1:
        mask.unsqueeze_(0)

    # Weighted Merge
    '根据pose的得分来决定每个目标所占的比例, 具体为该得分占总得分的比例'
    masked_scores = cluster_scores.mul(mask.float().unsqueeze(-1))
    normed_scores = masked_scores / torch.sum(masked_scores, dim=0)
	
	'根据计算得到的比例做加权和, 得到最终的pose及其得分'
    final_pose = torch.mul(cluster_preds, normed_scores.repeat(1, 1, 2)).sum(dim=0)
    final_score = torch.mul(masked_scores, normed_scores).sum(dim=0)
    return final_pose, final_score
final_result.append({
    'keypoints': merge_pose - 0.3,
    'kp_score': merge_score,
    'proposal_score': torch.mean(merge_score) + bbox_scores_pick[j] + 1.25 * max(merge_score)
})

keypoints 	[17,2]
kp_score	[17,1]
proposal_score [1]

3. Alphapose 网络结构

3.1 总流程

在这里插入图片描述

  1. SEResnet 作为 backbone 提取特征
  2. nn.PixelShuffle(2) 提升分辨率
  3. 经过两个 DUC 模块进一步提升分辨率
  4. 通过一个卷积得到输出
  5. 此时获得的输出如图所示 out[6,33,80,64] 有33个关键点,通过 out.narrow(1, 0, 17) 获取前17个关键点作为最终的输出 hm[6,17,80,64]

3.2 DUC 模块

  1. 2个DUC模块结构相同,都是先用卷积升维,再用一个 nn.PixelShuffle(2) 提升分辨率
  2. 图中以 DUC1 模块的参数为例进行绘制

3.3 PixelShuffle 操作

import torch
import torch.nn as nn

input_tensor = torch.arange(1, 17).view(1, 16, 1, 1).float()
pixel_shuffle = nn.PixelShuffle(2)
output_tensor = pixel_shuffle(input_tensor)

print(output_tensor)

>>>
tensor([[[[ 1.,  2.],
          [ 3.,  4.]],

         [[ 5.,  6.],
          [ 7.,  8.]],

         [[ 9., 10.],
          [11., 12.]],

         [[13., 14.],
          [15., 16.]]]])

3.4 SEResnet 框架

在这里插入图片描述
  图中省略了 batchsize 维度,主要分为4层,分别相对原图下采样4、8、16、32倍

3.5 SEResnet 细节

在这里插入图片描述
  仿照代码,把这4个由 Bottleneck_SEBottleneck 构成的层级记作 l a y e r 1 ∼ 4 mathrm{layer1sim4} layer14,图中为 l a y e r 1 mathrm{layer1} layer1 的数据。

  每个层中两种 Bottleneck 都会通过三个卷积层,先把特征维度控制为输出特征维度的 1/4,第二个保持不变,第三个达到输出特征维度,再以第二层为例:
l a y e r 2 : mathrm{layer2:} layer2:
   B o t t l e n e c k _ S E : 256 → 128 → 128 → 512 mathrm{Bottleneck_ SE:256to128to128to512} Bottleneck_SE:256128128512
   B o t t l e n e c k : 512 → 128 → 128 → 512 mathrm{Bottleneck:512to128to128to512} Bottleneck:512128128512