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)。

函数的主要步骤如下:

  1. 使用非极大值抑制(NMS)算法对热图进行处理,得到每个关节点类型的关节点列表(joint_list_per_joint_type)。
  2. 将关节点列表整合为一个数组(joint_list),其中每个关节点由坐标和关节点类型组成。
  3. 如果关节点列表不为空,则进行下一步处理。
  4. 对关节点列表进行扩展,添加一个额外的维度,用于后续处理。
  5. 使用插值方法(cv2.INTER_NEAREST)将向量场和热图上采样,以匹配原始图像尺寸。
  6. 使用 pafprocess 库处理关节点列表、热图和向量场,生成姿势估计结果。
  7. 遍历检测到的每个人体姿势,并为每个关节点创建一个 BodyPart 对象。
  8. 将关节点和姿势得分添加到人体对象中。
  9. 将检测到的人体对象添加到人体姿势列表中。
  10. 返回人体姿势列表。

你可能感兴趣的:(人工智能,计算机视觉,神经网络,深度学习)