You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
287 lines
11 KiB
287 lines
11 KiB
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. |
|
# |
|
# Licensed under the Apache License, Version 2.0 (the "License"); |
|
# you may not use this file except in compliance with the License. |
|
# You may obtain a copy of the License at |
|
# |
|
# http://www.apache.org/licenses/LICENSE-2.0 |
|
# |
|
# Unless required by applicable law or agreed to in writing, software |
|
# distributed under the License is distributed on an "AS IS" BASIS, |
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|
# See the License for the specific language governing permissions and |
|
# limitations under the License. |
|
|
|
from __future__ import absolute_import |
|
from __future__ import division |
|
from __future__ import print_function |
|
|
|
from scipy.optimize import linear_sum_assignment |
|
from collections import abc, defaultdict |
|
import numpy as np |
|
import paddle |
|
|
|
from paddlers_slim.models.ppdet.core.workspace import register, create, serializable |
|
from .meta_arch import BaseArch |
|
from .. import layers as L |
|
from ..keypoint_utils import transpred |
|
|
|
__all__ = ['HigherHRNet'] |
|
|
|
|
|
@register |
|
class HigherHRNet(BaseArch): |
|
__category__ = 'architecture' |
|
|
|
def __init__(self, |
|
backbone='HRNet', |
|
hrhrnet_head='HrHRNetHead', |
|
post_process='HrHRNetPostProcess', |
|
eval_flip=True, |
|
flip_perm=None, |
|
max_num_people=30): |
|
""" |
|
HigherHRNet network, see https://arxiv.org/abs/1908.10357; |
|
HigherHRNet+swahr, see https://arxiv.org/abs/2012.15175 |
|
|
|
Args: |
|
backbone (nn.Layer): backbone instance |
|
hrhrnet_head (nn.Layer): keypoint_head instance |
|
bbox_post_process (object): `BBoxPostProcess` instance |
|
""" |
|
super(HigherHRNet, self).__init__() |
|
self.backbone = backbone |
|
self.hrhrnet_head = hrhrnet_head |
|
self.post_process = post_process |
|
self.flip = eval_flip |
|
self.flip_perm = paddle.to_tensor(flip_perm) |
|
self.deploy = False |
|
self.interpolate = L.Upsample(2, mode='bilinear') |
|
self.pool = L.MaxPool(5, 1, 2) |
|
self.max_num_people = max_num_people |
|
|
|
@classmethod |
|
def from_config(cls, cfg, *args, **kwargs): |
|
# backbone |
|
backbone = create(cfg['backbone']) |
|
# head |
|
kwargs = {'input_shape': backbone.out_shape} |
|
hrhrnet_head = create(cfg['hrhrnet_head'], **kwargs) |
|
post_process = create(cfg['post_process']) |
|
|
|
return { |
|
'backbone': backbone, |
|
"hrhrnet_head": hrhrnet_head, |
|
"post_process": post_process, |
|
} |
|
|
|
def _forward(self): |
|
if self.flip and not self.training and not self.deploy: |
|
self.inputs['image'] = paddle.concat( |
|
(self.inputs['image'], paddle.flip(self.inputs['image'], [3]))) |
|
body_feats = self.backbone(self.inputs) |
|
|
|
if self.training: |
|
return self.hrhrnet_head(body_feats, self.inputs) |
|
else: |
|
outputs = self.hrhrnet_head(body_feats) |
|
|
|
if self.flip and not self.deploy: |
|
outputs = [paddle.split(o, 2) for o in outputs] |
|
output_rflip = [ |
|
paddle.flip(paddle.gather(o[1], self.flip_perm, 1), [3]) |
|
for o in outputs |
|
] |
|
output1 = [o[0] for o in outputs] |
|
heatmap = (output1[0] + output_rflip[0]) / 2. |
|
tagmaps = [output1[1], output_rflip[1]] |
|
outputs = [heatmap] + tagmaps |
|
outputs = self.get_topk(outputs) |
|
|
|
if self.deploy: |
|
return outputs |
|
|
|
res_lst = [] |
|
h = self.inputs['im_shape'][0, 0].numpy().item() |
|
w = self.inputs['im_shape'][0, 1].numpy().item() |
|
kpts, scores = self.post_process(*outputs, h, w) |
|
res_lst.append([kpts, scores]) |
|
return res_lst |
|
|
|
def get_loss(self): |
|
return self._forward() |
|
|
|
def get_pred(self): |
|
outputs = {} |
|
res_lst = self._forward() |
|
outputs['keypoint'] = res_lst |
|
return outputs |
|
|
|
def get_topk(self, outputs): |
|
# resize to image size |
|
outputs = [self.interpolate(x) for x in outputs] |
|
if len(outputs) == 3: |
|
tagmap = paddle.concat( |
|
(outputs[1].unsqueeze(4), outputs[2].unsqueeze(4)), axis=4) |
|
else: |
|
tagmap = outputs[1].unsqueeze(4) |
|
|
|
heatmap = outputs[0] |
|
N, J = 1, self.hrhrnet_head.num_joints |
|
heatmap_maxpool = self.pool(heatmap) |
|
# topk |
|
maxmap = heatmap * (heatmap == heatmap_maxpool) |
|
maxmap = maxmap.reshape([N, J, -1]) |
|
heat_k, inds_k = maxmap.topk(self.max_num_people, axis=2) |
|
|
|
outputs = [heatmap, tagmap, heat_k, inds_k] |
|
return outputs |
|
|
|
|
|
@register |
|
@serializable |
|
class HrHRNetPostProcess(object): |
|
''' |
|
HrHRNet postprocess contain: |
|
1) get topk keypoints in the output heatmap |
|
2) sample the tagmap's value corresponding to each of the topk coordinate |
|
3) match different joints to combine to some people with Hungary algorithm |
|
4) adjust the coordinate by +-0.25 to decrease error std |
|
5) salvage missing joints by check positivity of heatmap - tagdiff_norm |
|
Args: |
|
max_num_people (int): max number of people support in postprocess |
|
heat_thresh (float): value of topk below this threshhold will be ignored |
|
tag_thresh (float): coord's value sampled in tagmap below this threshold belong to same people for init |
|
|
|
inputs(list[heatmap]): the output list of model, [heatmap, heatmap_maxpool, tagmap], heatmap_maxpool used to get topk |
|
original_height, original_width (float): the original image size |
|
''' |
|
|
|
def __init__(self, max_num_people=30, heat_thresh=0.1, tag_thresh=1.): |
|
self.max_num_people = max_num_people |
|
self.heat_thresh = heat_thresh |
|
self.tag_thresh = tag_thresh |
|
|
|
def lerp(self, j, y, x, heatmap): |
|
H, W = heatmap.shape[-2:] |
|
left = np.clip(x - 1, 0, W - 1) |
|
right = np.clip(x + 1, 0, W - 1) |
|
up = np.clip(y - 1, 0, H - 1) |
|
down = np.clip(y + 1, 0, H - 1) |
|
offset_y = np.where(heatmap[j, down, x] > heatmap[j, up, x], 0.25, |
|
-0.25) |
|
offset_x = np.where(heatmap[j, y, right] > heatmap[j, y, left], 0.25, |
|
-0.25) |
|
return offset_y + 0.5, offset_x + 0.5 |
|
|
|
def __call__(self, heatmap, tagmap, heat_k, inds_k, original_height, |
|
original_width): |
|
|
|
N, J, H, W = heatmap.shape |
|
assert N == 1, "only support batch size 1" |
|
heatmap = heatmap[0].cpu().detach().numpy() |
|
tagmap = tagmap[0].cpu().detach().numpy() |
|
heats = heat_k[0].cpu().detach().numpy() |
|
inds_np = inds_k[0].cpu().detach().numpy() |
|
y = inds_np // W |
|
x = inds_np % W |
|
tags = tagmap[np.arange(J)[None, :].repeat(self.max_num_people), |
|
y.flatten(), x.flatten()].reshape(J, -1, tagmap.shape[-1]) |
|
coords = np.stack((y, x), axis=2) |
|
# threshold |
|
mask = heats > self.heat_thresh |
|
# cluster |
|
cluster = defaultdict(lambda: { |
|
'coords': np.zeros((J, 2), dtype=np.float32), |
|
'scores': np.zeros(J, dtype=np.float32), |
|
'tags': [] |
|
}) |
|
for jid, m in enumerate(mask): |
|
num_valid = m.sum() |
|
if num_valid == 0: |
|
continue |
|
valid_inds = np.where(m)[0] |
|
valid_tags = tags[jid, m, :] |
|
if len(cluster) == 0: # initialize |
|
for i in valid_inds: |
|
tag = tags[jid, i] |
|
key = tag[0] |
|
cluster[key]['tags'].append(tag) |
|
cluster[key]['scores'][jid] = heats[jid, i] |
|
cluster[key]['coords'][jid] = coords[jid, i] |
|
continue |
|
candidates = list(cluster.keys())[:self.max_num_people] |
|
centroids = [ |
|
np.mean( |
|
cluster[k]['tags'], axis=0) for k in candidates |
|
] |
|
num_clusters = len(centroids) |
|
# shape is (num_valid, num_clusters, tag_dim) |
|
dist = valid_tags[:, None, :] - np.array(centroids)[None, ...] |
|
l2_dist = np.linalg.norm(dist, ord=2, axis=2) |
|
# modulate dist with heat value, see `use_detection_val` |
|
cost = np.round(l2_dist) * 100 - heats[jid, m, None] |
|
# pad the cost matrix, otherwise new pose are ignored |
|
if num_valid > num_clusters: |
|
cost = np.pad(cost, ((0, 0), (0, num_valid - num_clusters)), |
|
'constant', |
|
constant_values=((0, 0), (0, 1e-10))) |
|
rows, cols = linear_sum_assignment(cost) |
|
for y, x in zip(rows, cols): |
|
tag = tags[jid, y] |
|
if y < num_valid and x < num_clusters and \ |
|
l2_dist[y, x] < self.tag_thresh: |
|
key = candidates[x] # merge to cluster |
|
else: |
|
key = tag[0] # initialize new cluster |
|
cluster[key]['tags'].append(tag) |
|
cluster[key]['scores'][jid] = heats[jid, y] |
|
cluster[key]['coords'][jid] = coords[jid, y] |
|
|
|
# shape is [k, J, 2] and [k, J] |
|
pose_tags = np.array([cluster[k]['tags'] for k in cluster]) |
|
pose_coords = np.array([cluster[k]['coords'] for k in cluster]) |
|
pose_scores = np.array([cluster[k]['scores'] for k in cluster]) |
|
valid = pose_scores > 0 |
|
|
|
pose_kpts = np.zeros((pose_scores.shape[0], J, 3), dtype=np.float32) |
|
if valid.sum() == 0: |
|
return pose_kpts, pose_kpts |
|
|
|
# refine coords |
|
valid_coords = pose_coords[valid].astype(np.int32) |
|
y = valid_coords[..., 0].flatten() |
|
x = valid_coords[..., 1].flatten() |
|
_, j = np.nonzero(valid) |
|
offsets = self.lerp(j, y, x, heatmap) |
|
pose_coords[valid, 0] += offsets[0] |
|
pose_coords[valid, 1] += offsets[1] |
|
|
|
# mean score before salvage |
|
mean_score = pose_scores.mean(axis=1) |
|
pose_kpts[valid, 2] = pose_scores[valid] |
|
|
|
# salvage missing joints |
|
if True: |
|
for pid, coords in enumerate(pose_coords): |
|
tag_mean = np.array(pose_tags[pid]).mean(axis=0) |
|
norm = np.sum((tagmap - tag_mean)**2, axis=3)**0.5 |
|
score = heatmap - np.round(norm) # (J, H, W) |
|
flat_score = score.reshape(J, -1) |
|
max_inds = np.argmax(flat_score, axis=1) |
|
max_scores = np.max(flat_score, axis=1) |
|
salvage_joints = (pose_scores[pid] == 0) & (max_scores > 0) |
|
if salvage_joints.sum() == 0: |
|
continue |
|
y = max_inds[salvage_joints] // W |
|
x = max_inds[salvage_joints] % W |
|
offsets = self.lerp(salvage_joints.nonzero()[0], y, x, heatmap) |
|
y = y.astype(np.float32) + offsets[0] |
|
x = x.astype(np.float32) + offsets[1] |
|
pose_coords[pid][salvage_joints, 0] = y |
|
pose_coords[pid][salvage_joints, 1] = x |
|
pose_kpts[pid][salvage_joints, 2] = max_scores[salvage_joints] |
|
pose_kpts[..., :2] = transpred(pose_coords[..., :2][..., ::-1], |
|
original_height, original_width, |
|
min(H, W)) |
|
return pose_kpts, mean_score
|
|
|