openpose保姆级教程代码详细解析——预测部分
目录
一：前言
二：预测 加载训练好的模型
三：加载预测模型的各个层
四：计算了关键点热图、人体姿态辅助图，图像的尺度变换因子
五:对图片进行缩放，计算缩放的比例因子
六：对图像进行预处理的函数
六：将处理后的图片输入加载好的训练好的模model(batch_var)
返回推断结果中的热图和向量场
七：将热图（heatmaps）和向量场（pafs）转换为姿势估计结果。
一：前言
OpenPose是一个基于深度学习的人体姿势估计库，它可以从图像或视频中准确地检测和估计人体的关键点和姿势信息。OpenPose的目标是将人体姿势估计变成一个实时、多人、准确的任务。——本节介绍预测部分
二：预测 加载训练好的模型
为了训练简单，使用了vgg19，可以自行修改不同的模型观测不同的预测结果
parser = argparse.ArgumentParser()
parser.add_argument('--cfg', help='experiment configure file name',
default='../experiments/vgg19_368x368_sgd.yaml', type=str)
parser.add_argument('--weight', type=str,
default='pose_model.pth')
parser.add_argument('opts',
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER) # 指定了参数后的其他命令行参数应该被收集为一个列表。这允许在命令行中传递多个选项。
args = parser.parse_args()
# update config file
# 这是一个函数调用，用于根据命令行参数更新配置文件。它可能是自定义的函数，其功能是将命令行参数应用于配置文件，以覆盖默认选项。
update_config(cfg, args)
model = get_model('vgg19')
model.load_state_dict(torch.load(args.weight)) # 加载pose_model.pth的权重
# 是将模型包装在torch.nn.DataParallel中，以实现在多个GPU上的并行训练。
model = torch.nn.DataParallel(model).cuda()
model.float()
model.eval() # 评估模式
test_image = 'ski.jpg'
oriImg = cv2.imread(test_image) # B,G,R order
# 获取了图像的最小尺寸（可能用于后续的处理或裁剪）。后面没用到
shape_dst = np.min(oriImg.shape[0:2])三：加载预测模型的各个层
class rtpose_model(nn.Module):
def __init__(self, model_dict):
super(rtpose_model, self).__init__()
self.model0 = model_dict['block0']
self.model1_1 = model_dict['block1_1']
self.model2_1 = model_dict['block2_1']
self.model3_1 = model_dict['block3_1']
self.model4_1 = model_dict['block4_1']
self.model5_1 = model_dict['block5_1']
self.model6_1 = model_dict['block6_1']
self.model1_2 = model_dict['block1_2']
self.model2_2 = model_dict['block2_2']
self.model3_2 = model_dict['block3_2']
self.model4_2 = model_dict['block4_2']
self.model5_2 = model_dict['block5_2']
self.model6_2 = model_dict['block6_2']
self._initialize_weights_norm()四：计算了关键点热图、人体姿态辅助图，图像的尺度变换因子
输入图片，模型，处理图片的模式‘rtpose’
# Get results of original image
with torch.no_grad():
# 计算了关键点热图（heatmap）、人体姿态辅助图（paf）以及图像的尺度变换因子（im_scale）。
paf, heatmap, im_scale = get_outputs(oriImg, model, 'rtpose')五:对图片进行缩放，计算缩放的比例因子
与yolov5的计算最少的填充非常类似，也是为了减小计算量吧
#像是yolov5的计算最少的填充类似
def crop_with_factor(im, dest_size=None, factor=32, is_ceil=True):
im_shape = im.shape
im_size_min = np.min(im_shape[0:2])
im_size_max = np.max(im_shape[0:2])
# im_scale = 1.
# if max_size is not None and im_size_min > max_size:
im_scale = float(dest_size) / im_size_min
im = cv2.resize(im, None, fx=im_scale, fy=im_scale)
h, w, c = im.shape
new_h = _factor_closest(h, factor=factor, is_ceil=is_ceil) # 变成8的倍数
new_w = _factor_closest(w, factor=factor, is_ceil=is_ceil)
im_croped = np.zeros([new_h, new_w, c], dtype=im.dtype)#补0
im_croped[0:h, 0:w, :] = im
return im_croped, im_scale, im.shape六：对图像进行预处理的函数
进行归一化方式不同，通道顺序不同，均值和标准差的定义不同，并且在 im_data 数组的第一个维度前插入一个新的维度
def get_outputs(img, model, preprocess):
"""Computes the averaged heatmap and paf for the given image
:param multiplier:
:param origImg: numpy array, the image being processed
:param model: pytorch model
:returns: numpy arrays, the averaged paf and heatmap
"""
inp_size = cfg.DATASET.IMAGE_SIZE
# padding im_croped 是裁剪并缩放后的图像。im_scale 是对图像进行缩放的比例因子。real_shape 是裁剪前的原始图像尺寸。
im_croped, im_scale, real_shape = im_transform.crop_with_factor(
img, inp_size, factor=cfg.MODEL.DOWNSAMPLE, is_ceil=True)
# 下面都是用于对图像进行预处理的函数，进行归一化方式不同，通道顺序不同，均值和标准差的定义不同
# 这些差异是基于不同的模型架构和训练设置而选择的预处理方式，以确保输入数据符合模型的期望格式和范围。
if preprocess == 'rtpose':
im_data = rtpose_preprocess(im_croped)
elif preprocess == 'vgg':
im_data = vgg_preprocess(im_croped)
elif preprocess == 'inception':
im_data = inception_preprocess(im_croped)
elif preprocess == 'ssd':
im_data = ssd_preprocess(im_croped)
# np.expand_dims(im_data, 0) 的目的是在 im_data 数组的第一个维度前插入一个新的维度。
batch_images= np.expand_dims(im_data, 0)六：将处理后的图片输入加载好的训练好的模model(batch_var)
进行forward 预测会通过模型的 forward 函数进行处理。
因为训练进行一次forward，然后还会进行backward，但是预测的时候只会进行forward
def forward(self, x):
# print(x.shape)
saved_for_loss = []
out1 = self.model0(x)#46*46的特征图
# print(out1.shape)
out1_1 = self.model1_1(out1) #PAF输出
# print(out1_1.shape)
out1_2 = self.model1_2(out1) #关键点输出
# print(out1_2.shape)
# cat操作把通道加起来 128+38+19 = 185 （185*46*46）
out2 = torch.cat([out1_1, out1_2, out1], 1)
# print(out2.shape)
# 每一个stage都要做损失
saved_for_loss.append(out1_1)
saved_for_loss.append(out1_2)
#out2每次上一层的输出。
out2_1 = self.model2_1(out2)
out2_2 = self.model2_2(out2)
out3 = torch.cat([out2_1, out2_2, out1], 1)
saved_for_loss.append(out2_1)
saved_for_loss.append(out2_2)
out3_1 = self.model3_1(out3)
out3_2 = self.model3_2(out3)
out4 = torch.cat([out3_1, out3_2, out1], 1)
saved_for_loss.append(out3_1)
saved_for_loss.append(out3_2)
out4_1 = self.model4_1(out4)
out4_2 = self.model4_2(out4)
out5 = torch.cat([out4_1, out4_2, out1], 1)
saved_for_loss.append(out4_1)
saved_for_loss.append(out4_2)
out5_1 = self.model5_1(out5)
out5_2 = self.model5_2(out5)
out6 = torch.cat([out5_1, out5_2, out1], 1)
saved_for_loss.append(out5_1)
saved_for_loss.append(out5_2)
out6_1 = self.model6_1(out6)
out6_2 = self.model6_2(out6)
saved_for_loss.append(out6_1)
saved_for_loss.append(out6_2)
# 这里应该就是对应的y的函数了，输入x根据训练好的权重w和b得到y
return (out6_1, out6_2), saved_for_loss
model = rtpose_model(models)
return model返回推断结果中的热图和向量场
# 从 output1 中提取出向量场（part affinity fields，简称 PAF），并将其转换为 NumPy 数组。类似地，
# 使用 output1.cpu().data.numpy().transpose(0, 2, 3, 1) 将其转换为形状为 (H, W, C) 的 PAF 数组。
# 从 output2 中提取出热图（heatmap），并将其转换为 NumPy 数组。这里使用 output2.cpu().data.numpy() 将 GPU 上的张量移回 CPU，
# 并使用 transpose(0, 2, 3, 1) 调整维度顺序，以得到形状为 (H, W, C) 的热图。
output1, output2 = predicted_outputs[-2], predicted_outputs[-1]
heatmap = output2.cpu().data.numpy().transpose(0, 2, 3, 1)[0]
paf = output1.cpu().data.numpy().transpose(0, 2, 3, 1)[0]
return paf, heatmap, im_scale七：将热图（heatmaps）和向量场（pafs）转换为姿势估计结果。
def paf_to_pose_cpp(heatmaps, pafs, config):
humans = []
joint_list_per_joint_type = NMS(heatmaps, upsampFactor=config.MODEL.DOWNSAMPLE, config=config)
joint_list = np.array(
[tuple(peak) + (joint_type,) for joint_type, joint_peaks in enumerate(joint_list_per_joint_type) for peak in
joint_peaks]).astype(np.float32)
if joint_list.shape[0] > 0:
joint_list = np.expand_dims(joint_list, 0)
paf_upsamp = cv2.resize(
pafs, None, fx=config.MODEL.DOWNSAMPLE, fy=config.MODEL.DOWNSAMPLE, interpolation=cv2.INTER_NEAREST)
heatmap_upsamp = cv2.resize(
heatmaps, None, fx=config.MODEL.DOWNSAMPLE, fy=config.MODEL.DOWNSAMPLE, interpolation=cv2.INTER_NEAREST)
pafprocess.process_paf(joint_list, heatmap_upsamp, paf_upsamp)
for human_id in range(pafprocess.get_num_humans()):
human = Human([])
is_added = False
for part_idx in range(config.MODEL.NUM_KEYPOINTS):
c_idx = int(pafprocess.get_part_cid(human_id, part_idx))
if c_idx < 0:
continue
is_added = True
human.body_parts[part_idx] = BodyPart(
'%d-%d' % (human_id, part_idx), part_idx,
float(pafprocess.get_part_x(c_idx)) / heatmap_upsamp.shape[1],
float(pafprocess.get_part_y(c_idx)) / heatmap_upsamp.shape[0],
pafprocess.get_part_score(c_idx)
)
if is_added:
score = pafprocess.get_score(human_id)
human.score = score
humans.append(human)
return humans
humans = paf_to_pose_cpp(heatmap, paf, cfg)
#将人体姿态信息绘制在原始图像上，并将结果保存为名为 'ski-output.jpg' 的图像文件。
out = draw_humans(oriImg, humans)
cv2.imwrite('ski-output.jpg', out)
函数的输入包括热图（heatmaps）、向量场（pafs）和配置参数（config），输出为检测到的人体姿势列表（humans）。
函数的主要步骤如下：
- 使用非极大值抑制（NMS）算法对热图进行处理，得到每个关节点类型的关节点列表（joint_list_per_joint_type）。
- 将关节点列表整合为一个数组（joint_list），其中每个关节点由坐标和关节点类型组成。
- 如果关节点列表不为空，则进行下一步处理。
- 对关节点列表进行扩展，添加一个额外的维度，用于后续处理。
- 使用插值方法（cv2.INTER_NEAREST）将向量场和热图上采样，以匹配原始图像尺寸。
- 使用 pafprocess 库处理关节点列表、热图和向量场，生成姿势估计结果。
- 遍历检测到的每个人体姿势，并为每个关节点创建一个 BodyPart 对象。
- 将关节点和姿势得分添加到人体对象中。
- 将检测到的人体对象添加到人体姿势列表中。
- 返回人体姿势列表。