Openpose 源码梳理

先看数据处理部分（datasets.py）

1.CocoKeypoints类中的getitem

取的index实际上就是一个随机值，也就是说从train数据集中随机取一张图片，因为论文中用的coco数据集，所以使用的都是coco官方提供的一些api，方便提取标注、图片信息。

def __getitem__(self, index):
    """
    Args:
        index (int): Index

    Returns:
        tuple: Tuple (image, target). target is the object returned by ``coco.loadAnns``.
    """
    image_id = self.ids[index]
    ann_ids = self.coco.getAnnIds(imgIds=image_id, catIds=self.cat_ids)
    anns = self.coco.loadAnns(ann_ids)
    anns = copy.deepcopy(anns)

    image_info = self.coco.loadImgs(image_id)[0]
    self.log.debug(image_info)
    with open(os.path.join(self.root, image_info['file_name']), 'rb') as f:
        image = Image.open(f).convert('RGB')

    meta_init = {
        'dataset_index': index,
        'image_id': image_id,
        'file_name': image_info['file_name'],
    }

    image, anns, meta = self.preprocess(image, anns, None)
         
    if isinstance(image, list):
        return self.multi_image_processing(image, anns, meta, meta_init)

    return self.single_image_processing(image, anns, meta, meta_init)

然后就是跳到multi_image_processing或是single_image_processing中，我们可以看出multi_image_processing也是基于single_image_processing

    def multi_image_processing(self, image_list, anns_list, meta_list, meta_init):
        return list(zip(*[
            self.single_image_processing(image, anns, meta, meta_init)
            for image, anns, meta in zip(image_list, anns_list, meta_list)
        ]))

所以我们只看single_image_processing部分就行。

2.single_image_processing

这段代码中最关键的是get_ground_truth，获取GT。

 def single_image_processing(self, image, anns, meta, meta_init):
        meta.update(meta_init)

        # transform image
        original_size = image.size
        image = self.image_transform(image)
        assert image.size(2) == original_size[0]
        assert image.size(1) == original_size[1]

        # mask valid
        valid_area = meta['valid_area']
        utils.mask_valid_area(image, valid_area)

        self.log.debug(meta)

        heatmaps, pafs = self.get_ground_truth(anns)
        
        heatmaps = torch.from_numpy(
            heatmaps.transpose((2, 0, 1)).astype(np.float32))
            
        pafs = torch.from_numpy(pafs.transpose((2, 0, 1)).astype(np.float32))       
        return image, heatmaps, pafs

3.get_ground_truth

这是整体代码，下面我们分开来说

    def get_ground_truth(self, anns):
    
        grid_y = int(self.input_y / self.stride)
        grid_x = int(self.input_x / self.stride)
        channels_heat = (self.HEATMAP_COUNT + 1)
        channels_paf = 2 * len(self.LIMB_IDS)
        heatmaps = np.zeros((int(grid_y), int(grid_x), channels_heat))
        pafs = np.zeros((int(grid_y), int(grid_x), channels_paf))

        keypoints = []
        for ann in anns:
            single_keypoints = np.array(ann['keypoints']).reshape(17,3)
            single_keypoints = self.add_neck(single_keypoints)
            keypoints.append(single_keypoints)
        keypoints = np.array(keypoints)
        keypoints = self.remove_illegal_joint(keypoints)

        # confidance maps for body parts
        for i in range(self.HEATMAP_COUNT):
            joints = [jo[i] for jo in keypoints]#每一种关节点
            for joint in joints:#遍历每一个点
                if joint[2] > 0.5:#1是标注被遮挡 2是标注且没被遮挡
                    center = joint[:2]#点坐标
                    gaussian_map = heatmaps[:, :, i]
                    heatmaps[:, :, i] = putGaussianMaps(
                        center, gaussian_map,
                        7.0, grid_y, grid_x, self.stride)
        # pafs
        for i, (k1, k2) in enumerate(self.LIMB_IDS):
            # limb
            count = np.zeros((int(grid_y), int(grid_x)), dtype=np.uint32) # 表示该位置是否被计算了多次（计算的数量）
            for joint in keypoints:
                if joint[k1, 2] > 0.5 and joint[k2, 2] > 0.5:
                    centerA = joint[k1, :2]
                    centerB = joint[k2, :2]
                    vec_map = pafs[:, :, 2 * i:2 * (i + 1)] #每一个躯干位置，选择x和y两个方向

                    pafs[:, :, 2 * i:2 * (i + 1)], count = putVecMaps(
                        centerA=centerA,
                        centerB=centerB,
                        accumulate_vec_map=vec_map,
                        count=count, grid_y=grid_y, grid_x=grid_x, stride=self.stride
                    )

        # background
        heatmaps[:, :, -1] = np.maximum(
            1 - np.max(heatmaps[:, :, :self.HEATMAP_COUNT], axis=2),
            0.
        )
        return heatmaps, pafs

我们的input_y和input_x=368，经过三次下采样也就是2^3=8，得到的特征图大小为46*46（368/8=46），即grid_x =grid_y=46
HEATMAP_COUNT=18，其中标注中有17个关键点，还有一个neck是通过左右肩膀自己手动计算出来的。
len(self.LIMB_IDS)=19，18个关键点有19种连接方式，（怎么连接人家已经给好了）。
所以这里面 channels_heat =19，加的1是表示背景。 channels_paf =38，乘2表示用x和y两个方向表示一个向量。
然后初始化heatmaps 和pafs ，分别初始化为(46,46,19)和(46,46,38)的全零矩阵。

		grid_y = int(self.input_y / self.stride)
        grid_x = int(self.input_x / self.stride)
        channels_heat = (self.HEATMAP_COUNT + 1)
        channels_paf = 2 * len(self.LIMB_IDS)
        heatmaps = np.zeros((int(grid_y), int(grid_x), channels_heat))
        pafs = np.zeros((int(grid_y), int(grid_x), channels_paf))

这里其实就是将标注中的keypoints 提取出来保存下来，并且还加上我们计算的neck 关键点。最后再过滤一遍keypoints 数据
（数据中可能有标注错误的情况，如果标注出来的位置x,y都超过了特征图大小范围（超出图片了），我们就把他过滤掉）

注：这里要提一句
在ann标注中，关键点(x,y,k)除了有x,y的坐标值以外还有一个k，k属于0,1,2中的一个。
k=0表示原图标注中就没有这个关键点（相对应的x,y值也是随机值，没有用，后面会过滤掉）
k=1表示关键点被遮挡，但是给出关键点的坐标
k=2表示正常的关键点，没有被遮挡
我们只用k=1和2的关键点。
所以这里的np.array(ann[‘keypoints’]).reshape(17,3)，中的17 表示标注中提供的17个关键点，3表示关键点对应的x,y,k

 		keypoints = []
        for ann in anns:
            single_keypoints = np.array(ann['keypoints']).reshape(17,3)
            single_keypoints = self.add_neck(single_keypoints)
            keypoints.append(single_keypoints)
        keypoints = np.array(keypoints)
        keypoints = self.remove_illegal_joint(keypoints)

然后开始构建heatmaps：

遍历第i个关键点：
从keypoints中取出每个人的关键点，将每个人第i个关键点存放在joints中。（这里的keypoints包含了所有人的所有关键点）
遍历第j个joints：
只取k=1或2的joint
保存joints的x，y
基于x，y生成高斯图（每次要基于之前的高斯图，因为重叠的部分要叠加）

 # confidance maps for body parts
        for i in range(self.HEATMAP_COUNT):
            joints = [jo[i] for jo in keypoints]#每一种关节点
            for joint in joints:#遍历每一个点
                if joint[2] > 0.5:#1是标注被遮挡 2是标注且没被遮挡
                    center = joint[:2]#点坐标
                    gaussian_map = heatmaps[:, :, i]
                    heatmaps[:, :, i] = putGaussianMaps(
                        center, gaussian_map,
                        7.0, grid_y, grid_x, self.stride)

然后开始构建pafs：
遍历第i个limb，同时拿到k1,k2信息（由k1,k2构成一个limb，LIMB_IDS里有指定都有哪些关键点相连）：
定义一个count (46,46)，用于记录每个点都包含几个limb信息（用于后面计算重叠部分，向量的重叠不是叠加，是取平均）
遍历第j个keypoints：（第j个人）
先看这个人的k1,k2关键点的k是否属于1或2
保存k1的x，y
保存k2的x，y
对于每个limb构建2个pafs
由k1的x，y和k2的x，y计算pafs

# pafs
        for i, (k1, k2) in enumerate(self.LIMB_IDS):
            # limb
            count = np.zeros((int(grid_y), int(grid_x)), dtype=np.uint32) # 表示该位置是否被计算了多次（计算的数量）
            for joint in keypoints:
                if joint[k1, 2] > 0.5 and joint[k2, 2] > 0.5:
                    centerA = joint[k1, :2]
                    centerB = joint[k2, :2]
                    vec_map = pafs[:, :, 2 * i:2 * (i + 1)] #每一个躯干位置，选择x和y两个方向

                    pafs[:, :, 2 * i:2 * (i + 1)], count = putVecMaps(
                        centerA=centerA,
                        centerB=centerB,
                        accumulate_vec_map=vec_map,
                        count=count, grid_y=grid_y, grid_x=grid_x, stride=self.stride
                    )

。。。。未完待续