• HigherHRNet 源码分析

    HigherHRNet 是一个 bottom-up 模型

    Bottom-up的人体姿势估计方法由于尺度变化的挑战，对于小人物的尺寸预测效果并不是太好。

    HigherHRNet 训练时采用多分辨率监督，简单来说就是输出多分辨率的heatmap

    实现也就是在HRNet的后面 添加一个反卷积分支，让其输出一个更大尺寸的heatmap

    关于HRNet 可以参考：

    HRNet 源码分析_那时那月那人的博客-CSDN博客

    下面分析 模型代码：基本上和HRNet基本一致，仅仅在后面添加一个反卷积分支

    1. class PoseHigherResolutionNet(nn.Module):
    2.  
    3.     def __init__(self, cfg, **kwargs):
    4.         pass
    5.  
    6.     def forward(self, x):
    7.         # 模型和HRNet基本一致 只是在最后添加一个返回更高分辨率的的输出
    8.         x = self.conv1(x)
    9.         x = self.bn1(x)
    10.         x = self.relu(x)
    11.         x = self.conv2(x)
    12.         x = self.bn2(x)
    13.         x = self.relu(x)
    14.         x = self.layer1(x)
    15.  
    16.         x_list = []
    17.         for i in range(self.stage2_cfg['NUM_BRANCHES']):
    18.             if self.transition1[i] is not None:
    19.                 x_list.append(self.transition1[i](x))
    20.             else:
    21.                 x_list.append(x)
    22.         y_list = self.stage2(x_list)
    23.  
    24.         x_list = []
    25.         for i in range(self.stage3_cfg['NUM_BRANCHES']):
    26.             if self.transition2[i] is not None:
    27.                 x_list.append(self.transition2[i](y_list[-1]))
    28.             else:
    29.                 x_list.append(y_list[i])
    30.         y_list = self.stage3(x_list)
    31.  
    32.         x_list = []
    33.         for i in range(self.stage4_cfg['NUM_BRANCHES']):
    34.             if self.transition3[i] is not None:
    35.                 x_list.append(self.transition3[i](y_list[-1]))
    36.             else:
    37.                 x_list.append(y_list[i])
    38.         y_list = self.stage4(x_list)
    39.  
    40.         final_outputs = []
    41.         x = y_list[0]
    42.         y = self.final_layers[0](x)
    43.         final_outputs.append(y)
    44.         # 在最后添加一个反卷积增大分辨率 如果需要更大的分比率可以加添加反卷积
    45.         # concat([features, heatmaps])
    46.         # 论文表面 添加一个反卷积在coco上达到最优
    47.         for i in range(self.num_deconvs):
    48.             if self.deconv_config.CAT_OUTPUT[i]:
    49.                 x = torch.cat((x, y), 1)
    50.  
    51.             x = self.deconv_layers[i](x)
    52.             y = self.final_layers[i + 1](x)
    53.             final_outputs.append(y)
    54.  
    55.         return final_outputs

    下面分析下heatmap label的生成代码，

    1. class CocoKeypoints(CocoDataset):
    2.     def __init__(self,
    3.                  cfg,
    4.                  dataset_name,
    5.                  remove_images_without_annotations,
    6.                  heatmap_generator,
    7.                  joints_generator,
    8.                  transforms=None):
    9.         super().__init__(cfg.DATASET.ROOT,
    10.                          dataset_name,
    11.                          cfg.DATASET.DATA_FORMAT)
    12.  
    13.         pass
    14.         self.heatmap_generator = heatmap_generator
    15.         self.joints_generator = joints_generator
    16.  
    17.     def __getitem__(self, idx):
    18.         # 调用COCODataset的__getitem()__得到 img, target
    19.         img, anno = super().__getitem__(idx)
    20.  
    21.         if img is None:
    22.             img_info = self.coco.loadImgs(self.ids[idx])[0]
    23.             img = np.zeros((3, img_info['height'], img_info['width']))
    24.  
    25.         mask = self.get_mask(anno, idx)
    26.  
    27.         anno = [
    28.             obj for obj in anno
    29.             if obj['iscrowd'] == 0 or obj['num_keypoints'] > 0
    30.         ]
    31.  
    32.         # TODO(bowen): to generate scale-aware sigma, modify `get_joints` to associate a sigma to each joint
    33.         joints = self.get_joints(anno)
    34.  
    35.         mask_list = [mask.copy() for _ in range(self.num_scales)]
    36.         joints_list = [joints.copy() for _ in range(self.num_scales)]
    37.         target_list = list()
    38.  
    39.         if self.transforms:
    40.             img, mask_list, joints_list = self.transforms(
    41.                 img, mask_list, joints_list
    42.             )
    43.  
    44.         for scale_id in range(self.num_scales):
    45.             # 对多尺寸生成heatmap
    46.             target_t = self.heatmap_generator[scale_id](joints_list[scale_id])
    47.             joints_t = self.joints_generator[scale_id](joints_list[scale_id])
    48.  
    49.             target_list.append(target_t.astype(np.float32))
    50.             mask_list[scale_id] = mask_list[scale_id].astype(np.float32)
    51.             joints_list[scale_id] = joints_t.astype(np.int32)
    52.  
    53.         return img, target_list, mask_list, joints_list
    1. class HeatmapGenerator():
    2.     def __init__(self, output_res, num_joints, sigma=-1):
    3.         # heatmap输出的尺寸
    4.         self.output_res = output_res
    5.         # 17
    6.         self.num_joints = num_joints
    7.         # w32_512_adam_lr1e-3.yaml   sigma = 2
    8.         if sigma < 0:
    9.             sigma = self.output_res/64
    10.         self.sigma = sigma
    11.         # 高斯核的大小
    12.         size = 6*sigma + 3
    13.         x = np.arange(0, size, 1, float)
    14.         y = x[:, np.newaxis]
    15.         x0, y0 = 3*sigma + 1, 3*sigma + 1
    16.         # 用2d高斯函数生成 label
    17.         self.g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
    18.  
    19.     def __call__(self, joints):
    20.         # heatmaps 尺寸
    21.         hms = np.zeros((self.num_joints, self.output_res, self.output_res),
    22.                        dtype=np.float32)
    23.         sigma = self.sigma
    24.         for p in joints:
    25.             for idx, pt in enumerate(p):
    26.                 if pt[2] > 0:
    27.                     # 对每个点用高斯函数来生成label
    28.                     x, y = int(pt[0]), int(pt[1])
    29.                     if x < 0 or y < 0 or \
    30.                        x >= self.output_res or y >= self.output_res:
    31.                         continue
    32.  
    33.                     ul = int(np.round(x - 3 * sigma - 1)), int(np.round(y - 3 * sigma - 1))
    34.                     br = int(np.round(x + 3 * sigma + 2)), int(np.round(y + 3 * sigma + 2))
    35.  
    36.                     c, d = max(0, -ul[0]), min(br[0], self.output_res) - ul[0]
    37.                     a, b = max(0, -ul[1]), min(br[1], self.output_res) - ul[1]
    38.  
    39.                     cc, dd = max(0, ul[0]), min(br[0], self.output_res)
    40.                     aa, bb = max(0, ul[1]), min(br[1], self.output_res)
    41.                     # 高斯函数赋值 
    42.                     hms[idx, aa:bb, cc:dd] = np.maximum(
    43.                         hms[idx, aa:bb, cc:dd], self.g[a:b, c:d])
    44.         return hms

    论文中提出一种 尺寸感知 （scale-aware）的训练方法 下面分析下具体的实现

    从 scale-ware 字面意思 尺寸感知，根据不同尺寸得到不同的高斯函数。从代码中可以看出根据关键点的可见度来进行判断。不太懂？可见度和尺寸有什么联系？

    1. class ScaleAwareHeatmapGenerator():
    2.     def __init__(self, output_res, num_joints):
    3.         self.output_res = output_res
    4.         self.num_joints = num_joints
    5.  
    6.     def get_gaussian_kernel(self, sigma):
    7.         size = 6*sigma + 3
    8.         x = np.arange(0, size, 1, float)
    9.         y = x[:, np.newaxis]
    10.         x0, y0 = 3*sigma + 1, 3*sigma + 1
    11.         g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
    12.         return g
    13.  
    14.     def __call__(self, joints):
    15.         hms = np.zeros((self.num_joints, self.output_res, self.output_res),
    16.                        dtype=np.float32)
    17.         for p in joints:
    18.             sigma = p[0, 3]
    19.             # 根据 可见性生成不同的高斯函数
    20.             # simga:
    21.             # 0: 未标注 1： 标注不可见 2： 标注可见
    22.             # 这不同的尺寸感知 和 可见性有关联？ 越小的的尺寸越不可见？
    23.             g = self.get_gaussian_kernel(sigma)
    24.             for idx, pt in enumerate(p):
    25.                 if pt[2] > 0:
    26.                     x, y = int(pt[0]), int(pt[1])
    27.                     if x < 0 or y < 0 or \
    28.                        x >= self.output_res or y >= self.output_res:
    29.                         continue
    30.  
    31.                     ul = int(np.round(x - 3 * sigma - 1)), int(np.round(y - 3 * sigma - 1))
    32.                     br = int(np.round(x + 3 * sigma + 2)), int(np.round(y + 3 * sigma + 2))
    33.  
    34.                     c, d = max(0, -ul[0]), min(br[0], self.output_res) - ul[0]
    35.                     a, b = max(0, -ul[1]), min(br[1], self.output_res) - ul[1]
    36.  
    37.                     cc, dd = max(0, ul[0]), min(br[0], self.output_res)
    38.                     aa, bb = max(0, ul[1]), min(br[1], self.output_res)
    39.                     hms[idx, aa:bb, cc:dd] = np.maximum(
    40.                         hms[idx, aa:bb, cc:dd], g[a:b, c:d])
    41.         return hms

    接下来分析损失函数

    一般bottom-up是先找点，然后在分组。openpose 采用的PAF 来进行划分。

    HigherHRNet使用 associate embedding 方法来进行分组的。

    如果不了解 associate embedding 请关注下面这篇论文

    Associative Embedding: End-to-End Learning for Joint Detection and Grouping

    https://arxiv.org/abs/1611.05424

     Lg损失函数又上面两部分组成。

    第一部分表示 同一个人所有关键点之间的损失函数，让所有关键点 embedding 先 reference embedding（平均 embedding）靠近。

    第二部分表示 不同人之间的损失，两着reference embedding 距离越大 损失函数越小 exp(-d**2)

    接下来分析下损失函数的代码实现

    1. class MultiLossFactory(nn.Module):
    2.     def __init__(self, cfg):
    3.         super().__init__()
    4.         # init check
    5.         self._init_check(cfg)
    6.  
    7.         self.num_joints = cfg.MODEL.NUM_JOINTS
    8.         self.num_stages = cfg.LOSS.NUM_STAGES
    9.  
    10.         # heatmap 损失 MSE 损失
    11.         # WITH_HEATMAPS_LOSS: [True, True]
    12.         # [HeatmapLoss(), HeatmapLoss()]
    13.         self.heatmaps_loss = \
    14.             nn.ModuleList(
    15.                 [
    16.                     HeatmapLoss()
    17.                     if with_heatmaps_loss else None
    18.                     for with_heatmaps_loss in cfg.LOSS.WITH_HEATMAPS_LOSS
    19.                 ]
    20.             )
    21.         self.heatmaps_loss_factor = cfg.LOSS.HEATMAPS_LOSS_FACTOR
    22.  
    23.         # associative embedding 损失
    24.         # WITH_AE_LOSS: [True, False]
    25.         # [AELoss(), None]
    26.         self.ae_loss = \
    27.             nn.ModuleList(
    28.                 [
    29.                     AELoss(cfg.LOSS.AE_LOSS_TYPE) if with_ae_loss else None
    30.                     for with_ae_loss in cfg.LOSS.WITH_AE_LOSS
    31.                 ]
    32.             )
    33.         self.push_loss_factor = cfg.LOSS.PUSH_LOSS_FACTOR
    34.         self.pull_loss_factor = cfg.LOSS.PULL_LOSS_FACTOR
    35.  
    36.     def forward(self, outputs, heatmaps, masks, joints):
    37.         # loss_factory(outputs, heatmaps, masks, joints)
    38.         # forward check
    39.         self._forward_check(outputs, heatmaps, masks, joints)
    40.  
    41.         heatmaps_losses = []
    42.         push_losses = []
    43.         pull_losses = []
    44.         # len(outputs) == 2
    45.         for idx in range(len(outputs)):
    46.             offset_feat = 0
    47.             # [HeatmapLoss(), HeatmapLoss()]
    48.             if self.heatmaps_loss[idx]:
    49.                 heatmaps_pred = outputs[idx][:, :self.num_joints]
    50.                 offset_feat = self.num_joints
    51.                 # 计算 heatmap 损失
    52.                 heatmaps_loss = self.heatmaps_loss[idx](
    53.                     heatmaps_pred, heatmaps[idx], masks[idx]
    54.                 )
    55.                 # self.heatmaps_loss_factor[idx] 损失函数权重系数 1
    56.                 heatmaps_loss = heatmaps_loss * self.heatmaps_loss_factor[idx]
    57.                 heatmaps_losses.append(heatmaps_loss)
    58.             else:
    59.                 heatmaps_losses.append(None)
    60.  
    61.             # Associative Embedding
    62.             # [AELoss(), None]
    63.             if self.ae_loss[idx]:
    64.                 tags_pred = outputs[idx][:, offset_feat:]
    65.                 batch_size = tags_pred.size()[0]
    66.                 tags_pred = tags_pred.contiguous().view(batch_size, -1, 1)
    67.                 # 计算 associative embedding loss
    68.                 push_loss, pull_loss = self.ae_loss[idx](
    69.                     tags_pred, joints[idx]
    70.                 )
    71.                 push_loss = push_loss * self.push_loss_factor[idx]
    72.                 pull_loss = pull_loss * self.pull_loss_factor[idx]
    73.  
    74.                 push_losses.append(push_loss)
    75.                 pull_losses.append(pull_loss)
    76.             else:
    77.                 push_losses.append(None)
    78.                 pull_losses.append(None)
    79.  
    80.         return heatmaps_losses, push_losses, pull_losses

    下面分析ae (associative embedding) 损失

    1. class AELoss(nn.Module):
    2.     def __init__(self, loss_type):
    3.         super().__init__()
    4.         self.loss_type = loss_type
    5.  
    6.     def singleTagLoss(self, pred_tag, joints):
    7.         """
    8.         associative embedding loss for one image
    9.         """
    10.         # pred_tag (17 * h * w, 1)
    11.         # joints (30, 17, 2)
    12.         tags = []
    13.         pull = 0
    14.         # 遍历人数
    15.         for joints_per_person in joints:
    16.             tmp = []
    17.             # 遍历所有关节点
    18.             for joint in joints_per_person:
    19.                 if joint[1] > 0:
    20.                     # joint 里面存的值 [关节点在heatmap的位置, 1]
    21.                     # 该关节点可见 加入tmp
    22.                     tmp.append(pred_tag[joint[0]])
    23.             # 如果没有可见关键点 continue
    24.             if len(tmp) == 0:
    25.                 continue
    26.             tmp = torch.stack(tmp)
    27.             tags.append(torch.mean(tmp, dim=0))
    28.             # 这里是 对单个人得损失 让每个关节点 embedding 向 reference embedding （均值靠拢）
    29.             pull = pull + torch.mean((tmp - tags[-1].expand_as(tmp))**2)
    30.  
    31.         num_tags = len(tags)
    32.         if num_tags == 0:
    33.             return make_input(torch.zeros(1).float()), \
    34.                 make_input(torch.zeros(1).float())
    35.         elif num_tags == 1:
    36.             return make_input(torch.zeros(1).float()), \
    37.                 pull/(num_tags)
    38.  
    39.         tags = torch.stack(tags)
    40.  
    41.         # 不同人之间得损失
    42.         size = (num_tags, num_tags)
    43.         A = tags.expand(*size)
    44.         B = A.permute(1, 0)
    45.  
    46.         diff = A - B
    47.  
    48.         if self.loss_type == 'exp':
    49.             # 根据损失函数 希望 不同人之间reference embedding的距离越来越远
    50.             diff = torch.pow(diff, 2)
    51.             push = torch.exp(-diff)
    52.             # 着里 -num_tags 可有可无 仅仅是添加一个margin
    53.             push = torch.sum(push) - num_tags
    54.         elif self.loss_type == 'max':
    55.             diff = 1 - torch.abs(diff)
    56.             push = torch.clamp(diff, min=0).sum() - num_tags
    57.         else:
    58.             raise ValueError('Unkown ae loss type')
    59.  
    60.         return push/((num_tags - 1) * num_tags) * 0.5, \
    61.             pull/(num_tags)
    62.  
    63.     def forward(self, tags, joints):
    64.         """
    65.         accumulate the tag loss for each image in the batch
    66.         """
    67.         # tags (B, 17 * h * w, 1)
    68.         # joints (B, 30, 17, 2)
    69.         pushes, pulls = [], []
    70.         joints = joints.cpu().data.numpy()
    71.         batch_size = tags.size(0)
    72.         for i in range(batch_size):
    73.             push, pull = self.singleTagLoss(tags[i], joints[i])
    74.             pushes.append(push)
    75.             pulls.append(pull)
    76.         return torch.stack(pushes), torch.stack(pulls)

    接下来分析inference代码  论文中指出使用一个

    Heatmap Aggregation for Inference
    We propose a heatmap aggregation strategy during inference. We use bilinear interpolation to upsample all the predicted heatmaps with different resolutions to the reso lution of the input image and average the heatmaps from all scales for fifinal prediction. This strategy is quite different from previous methods [ 3 , 30 , 33 ] which only use heatmaps from a single scale or single stage for prediction.

    简单来说 使用多尺度的heatmap来进行预测，对于不同尺寸的heatmap采用双向性插值的上采用来进行处理。和其他的模型处理方法不同，其他模型处理方法一般只取最后一个stage的最大尺寸heatmap来进行预测。HigherHRNet之所以使用不同尺寸的heatmap来进行预测，不同的heatmap尺寸对不同尺度的关键点起到的作用也不同。低分辨率的heatmap容易漏掉小尺寸人的关键点。但是小尺寸人的关键点可以在高分辨率的heatmap中恢复。

    下面我们分析下inference代码  这里不同尺度的融合是采用取平均值的方式实现

    1. def get_multi_stage_outputs(
    2.         cfg, model, image, with_flip=False,
    3.         project2image=False, size_projected=None
    4. ):
    5.     # outputs = []
    6.     heatmaps_avg = 0
    7.     num_heatmaps = 0
    8.     heatmaps = []
    9.     tags = []
    10.  
    11.     outputs = model(image)
    12.     for i, output in enumerate(outputs):
    13.         if len(outputs) > 1 and i != len(outputs) - 1:
    14.             # 对小尺寸的heatmap用双线性插值放大
    15.             output = torch.nn.functional.interpolate(
    16.                 output,
    17.                 size=(outputs[-1].size(2), outputs[-1].size(3)),
    18.                 mode='bilinear',
    19.                 align_corners=False
    20.             )
    21.  
    22.         offset_feat = cfg.DATASET.NUM_JOINTS \
    23.             if cfg.LOSS.WITH_HEATMAPS_LOSS[i] else 0
    24.  
    25.         if cfg.LOSS.WITH_HEATMAPS_LOSS[i] and cfg.TEST.WITH_HEATMAPS[i]:
    26.             # 这里进行累加 后面计算平均值
    27.             heatmaps_avg += output[:, :cfg.DATASET.NUM_JOINTS]
    28.             num_heatmaps += 1
    29.  
    30.         if cfg.LOSS.WITH_AE_LOSS[i] and cfg.TEST.WITH_AE[i]:
    31.             # 这里是分组信息
    32.             tags.append(output[:, offset_feat:])
    33.  
    34.     if num_heatmaps > 0:
    35.         heatmaps.append(heatmaps_avg/num_heatmaps)

    HigherHRNet 主要代码分析完。

  • 相关阅读:
    APP自动化测试 ---- Appium介绍及运行原理
    Leetcode 1331. 数组序号转换
    有关于脉动调查的这些问题你都知道吗
    保姆级银河麒麟V10高级服务器离线安装mysql5.7数据库
    AttributeError: Can only use .dt accessor with datetimelike values
    enumerate内置函数的使用
    匿名内部类的概念和使用详解
    php 计算工作时间 排除节假日可设置补班
    C++之5|组合与继承
    淘宝/天猫邻家好货 API 返回值说明
  • 原文地址：https://blog.csdn.net/xiaoxu1025/article/details/127850559