tools/accuracy_checker/accuracy_checker/adapters/pose_estimation.py

   1 """
   2 Copyright (c) 2019 Intel Corporation
   3
   4 Licensed under the Apache License, Version 2.0 (the "License");
   5 you may not use this file except in compliance with the License.
   6 You may obtain a copy of the License at
   7
   8       http://www.apache.org/licenses/LICENSE-2.0
   9
  10 Unless required by applicable law or agreed to in writing, software
  11 distributed under the License is distributed on an "AS IS" BASIS,
  12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 See the License for the specific language governing permissions and
  14 limitations under the License.
  15 """
  16
  17 import math
  18 from operator import itemgetter
  19
  20 import cv2
  21 import numpy as np
  22
  23 from ..adapters import Adapter
  24 from ..config import ConfigValidator, StringField
  25 from ..representation import PoseEstimationPrediction
  26
  27
  28 class HumanPoseAdapterConfig(ConfigValidator):
  29     type = StringField()
  30     part_affinity_fields_out = StringField()
  31     keypoints_heatmap_out = StringField()
  32
  33
  34 class HumanPoseAdapter(Adapter):
  35     __provider__ = 'human_pose_estimation'
  36
  37     limb_seq = [
  38         [2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], [10, 11], [2, 12], [12, 13],
  39         [13, 14], [2, 1], [1, 15], [15, 17], [1, 16], [16, 18], [3, 17], [6, 18]
  40     ]
  41     map_idx = [
  42         [31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], [23, 24], [25, 26],
  43         [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], [55, 56], [37, 38], [45, 46]
  44     ]
  45
  46     def validate_config(self):
  47         human_pose_estimation_config = HumanPoseAdapterConfig('HumanPose_Config')
  48         human_pose_estimation_config.validate(self.launcher_config)
  49
  50     def configure(self):
  51         self.part_affinity_fields = self.launcher_config['part_affinity_fields_out']
  52         self.keypoints_heatmap = self.launcher_config['keypoints_heatmap_out']
  53
  54     def process(self, raw, identifiers=None, frame_meta=None):
  55         result = []
  56         raw_outputs = self._extract_predictions(raw, frame_meta)
  57         raw_output = zip(
  58             identifiers, raw_outputs[self.keypoints_heatmap],
  59             raw_outputs[self.part_affinity_fields], frame_meta
  60         )
  61         for identifier, heatmap, paf, meta in raw_output:
  62             height, width, _ = meta['image_size']
  63             heatmap_avg = np.zeros((height, width, 19), dtype=np.float32)
  64             paf_avg = np.zeros((height, width, 38), dtype=np.float32)
  65             pad = meta.get('padding', [0, 0, 0, 0])
  66             heatmap = np.transpose(np.squeeze(heatmap), (1, 2, 0))
  67             heatmap = cv2.resize(heatmap, (0, 0), fx=8, fy=8, interpolation=cv2.INTER_CUBIC)
  68             heatmap = heatmap[pad[0]:heatmap.shape[0] - pad[2], pad[1]:heatmap.shape[1] - pad[3]:, :]
  69             heatmap = cv2.resize(heatmap, (width, height), interpolation=cv2.INTER_CUBIC)
  70             heatmap_avg = heatmap_avg + heatmap
  71
  72             paf = np.transpose(np.squeeze(paf), (1, 2, 0))
  73             paf = cv2.resize(paf, (0, 0), fx=8, fy=8, interpolation=cv2.INTER_CUBIC)
  74             paf = paf[pad[0]:paf.shape[0] - pad[2], pad[1]:paf.shape[1] - pad[3], :]
  75             paf = cv2.resize(paf, (width, height), interpolation=cv2.INTER_CUBIC)
  76             paf_avg = paf_avg + paf
  77
  78             peak_counter = 0
  79             all_peaks = []
  80             for part in range(0, 18):  # 19th for bg
  81                 peak_counter += self.find_peaks(heatmap_avg[:, :, part], all_peaks, peak_counter)
  82
  83             subset, candidate = self.group_peaks(all_peaks, paf_avg)
  84             result.append(PoseEstimationPrediction(identifier, *self.get_poses(subset, candidate)))
  85
  86         return result
  87
  88     @staticmethod
  89     def find_peaks(heatmap, all_peaks, prev_peak_counter):
  90         heatmap[heatmap < 0.1] = 0
  91         map_aug = np.zeros((heatmap.shape[0] + 2, heatmap.shape[1] + 2))
  92         map_left = np.zeros(map_aug.shape)
  93         map_right = np.zeros(map_aug.shape)
  94         map_up = np.zeros(map_aug.shape)
  95         map_down = np.zeros(map_aug.shape)
  96
  97         map_aug[1:map_aug.shape[0] - 1, 1:map_aug.shape[1] - 1] = heatmap
  98         map_left[1:map_aug.shape[0] - 1, :map_aug.shape[1] - 2] = heatmap
  99         map_right[1:map_aug.shape[0] - 1, 2:map_aug.shape[1]] = heatmap
 100         map_up[:map_aug.shape[0] - 2, 1:map_aug.shape[1] - 1] = heatmap
 101         map_down[2:map_aug.shape[0], 1:map_aug.shape[1] - 1] = heatmap
 102
 103         peaks_binary = (map_aug > map_left) & (map_aug > map_right) & (map_aug > map_up) & (map_aug > map_down)
 104         peaks_binary = peaks_binary[1:map_aug.shape[0] - 1, 1:map_aug.shape[1] - 1]
 105         peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]))
 106         peaks = sorted(peaks, key=itemgetter(0))  # same order with matlab
 107
 108         flag = np.ones(len(peaks), np.uint8)
 109         peaks_with_score_and_id = []
 110         peak_counter = 0
 111         for i, _ in enumerate(peaks):
 112             if flag[i] != 1:
 113                 continue
 114             for j in range(i + 1, len(peaks)):
 115                 if math.sqrt((peaks[i][0] - peaks[j][0]) ** 2 + (peaks[i][1] - peaks[j][1]) ** 2) < 6:
 116                     flag[j] = 0
 117             peak_id = peak_counter + prev_peak_counter
 118             peak_counter += 1
 119             peaks_with_score_and_id.append([peaks[i][0], peaks[i][1], heatmap[peaks[i][1], peaks[i][0]], peak_id])
 120         all_peaks.append(peaks_with_score_and_id)
 121
 122         return peak_counter
 123
 124     @staticmethod
 125     def _add_pose_single_candidate(subset, candidate, idx_joint, kpt_num=20):
 126         for joint in candidate:
 127             num = 0
 128             for subset_j in subset:  # check if already in some pose, was added as a part of another limb
 129                 if subset_j[idx_joint] == joint[3]:
 130                     num += 1
 131                     continue
 132             if num == 0:
 133                 person_keypoints = np.ones(kpt_num) * -1
 134                 person_keypoints[idx_joint] = joint[3]  # joint idx
 135                 person_keypoints[-1] = 1  # n joints in pose
 136                 person_keypoints[-2] = joint[2]  # pose score
 137                 subset.append(person_keypoints)
 138
 139         return subset
 140
 141     @staticmethod
 142     def _filter_subset(subset):
 143         filtered_subset = []
 144         for subset_element in subset:
 145             if subset_element[-1] < 3 or (subset_element[-2] / subset_element[-1] < 0.2):
 146                 continue
 147             filtered_subset.append(subset_element)
 148
 149         return np.asarray(filtered_subset)
 150
 151     @staticmethod
 152     def _add_pose_both_candidates(subset, temp, index_a, index_b, candidates, kpt_num=20):
 153         for i, temp_i in enumerate(temp):
 154             num = 0
 155             for j, subset_j in enumerate(subset):
 156                 if subset_j[index_a] == temp_i[0]:
 157                     subset[j][index_b] = temp[i][1]
 158                     num += 1
 159                     subset[j][-1] += 1
 160                     subset[j][-2] += candidates[temp_i[1], 2] + temp_i[2]
 161             if num == 0:
 162                 person_keypoints = np.ones(kpt_num) * -1
 163                 person_keypoints[index_a] = temp[i][0]
 164                 person_keypoints[index_b] = temp[i][1]
 165                 person_keypoints[-1] = 2
 166                 person_keypoints[-2] = np.sum(candidates[temp_i[0:2], 2]) + temp_i[2]
 167                 subset.append(person_keypoints)
 168
 169         return subset
 170
 171     @staticmethod
 172     def _copy_temperature_to_subset(subset, temp, index_a, index_b):
 173         for _, temp_i in enumerate(temp):
 174             for j, subset_j in enumerate(subset):
 175                 check_subset_a = subset_j[index_a] == temp_i[0] and subset_j[index_b] == -1
 176                 check_subset_b = subset_j[index_b] == temp_i[1] and subset_j[index_a] == -1
 177                 if check_subset_a:
 178                     subset[j][index_b] = temp_i[1]
 179                     continue
 180                 if check_subset_b:
 181                     subset[j][index_a] = temp_i[0]
 182
 183         return subset
 184
 185     @staticmethod
 186     def _get_temperature(cand_a_, cand_b_, score_mid, pafs, threshold=0.05):
 187         temp_ = []
 188         for index_a_, cand_a_element in enumerate(cand_a_):
 189             for index_b_, cand_b_element in enumerate(cand_b_):
 190                 mid_point = [(
 191                     int(round((cand_a_element[0] + cand_b_element[0]) * 0.5)),
 192                     int(round((cand_a_element[1] + cand_b_element[1]) * 0.5))
 193                 )] * 2
 194                 vec = [cand_b_element[0] - cand_a_element[0], cand_b_element[1] - cand_a_element[1]]
 195                 norm_vec = math.sqrt(vec[0] ** 2 + vec[1] ** 2)
 196                 if norm_vec == 0:
 197                     continue
 198                 vec[0] /= norm_vec
 199                 vec[1] /= norm_vec
 200                 score_mid_a = score_mid[mid_point[0][1], mid_point[0][0], 0]
 201                 score_mid_b = score_mid[mid_point[1][1], mid_point[1][0], 1]
 202                 score = vec[0] * score_mid_a + vec[1] * score_mid_b
 203
 204                 height_n = pafs.shape[0] // 2
 205                 suc_ratio = 0
 206                 mid_score = 0
 207                 mid_num = 10  # n points for integral over paf
 208
 209                 if score > -100:
 210                     p_sum = 0
 211                     p_count = 0
 212
 213                     x = np.linspace(cand_a_element[0], cand_b_element[0], mid_num)
 214                     y = np.linspace(cand_a_element[1], cand_b_element[1], mid_num)
 215                     for point_idx in range(0, mid_num):
 216                         px = int(round(x[point_idx]))
 217                         py = int(round(y[point_idx]))
 218                         pred = score_mid[py, px, 0:2]
 219                         score = vec[0] * pred[0] + vec[1] * pred[1]
 220                         if score > threshold:
 221                             p_sum += score
 222                             p_count += 1
 223                     suc_ratio = p_count / mid_num
 224                     ratio = 0
 225                     if p_count > 0:
 226                         ratio = p_sum / p_count
 227                     mid_score = ratio + min(height_n / norm_vec - 1, 0)
 228                 if mid_score > 0 and suc_ratio > 0.8:
 229                     score = mid_score
 230                     score_all = score + cand_a_element[2] + cand_b_element[2]
 231                     temp_.append([index_a_, index_b_, score, score_all])
 232         if temp_:
 233             temp_ = sorted(temp_, key=itemgetter(2), reverse=True)
 234
 235         return temp_
 236
 237     def _get_connections(self, cand_a, cand_b, score_mid, pafs, thresh):
 238         temp_ = self._get_temperature(cand_a, cand_b, score_mid, pafs, thresh)
 239         num_limbs = min(len(cand_a), len(cand_b))
 240         cnt = 0
 241         occur_a = np.zeros(len(cand_a), dtype=np.int32)
 242         occur_b = np.zeros(len(cand_b), dtype=np.int32)
 243         connections = []
 244         for row_temp in temp_:
 245             if cnt == num_limbs:
 246                 break
 247             i, j, score = row_temp[0:3]
 248             if occur_a[i] == 0 and occur_b[j] == 0:
 249                 connections.append([cand_a[i][3], cand_b[j][3], score])
 250                 cnt += 1
 251                 occur_a[i] = 1
 252                 occur_b[j] = 1
 253         return connections
 254
 255     def group_peaks(self, peaks, pafs, kpt_num=20, threshold=0.05):
 256         subset = []
 257         candidates = np.array([item for sublist in peaks for item in sublist])
 258         for keypoint_id, maped_keypoints in enumerate(self.map_idx):
 259             score_mid = pafs[:, :, [x - 19 for x in maped_keypoints]]
 260             candidate_a = peaks[self.limb_seq[keypoint_id][0] - 1]
 261             candidate_b = peaks[self.limb_seq[keypoint_id][1] - 1]
 262             idx_joint_a = self.limb_seq[keypoint_id][0] - 1
 263             idx_joint_b = self.limb_seq[keypoint_id][1] - 1
 264
 265             if not candidate_a and not candidate_b:  # no such limb
 266                 continue
 267             if not candidate_a:  # limb has just B joint
 268                 subset = self._add_pose_single_candidate(subset, candidate_b, idx_joint_b, kpt_num)
 269                 continue
 270             if not candidate_b:  # limb has just A joint
 271                 subset = self._add_pose_single_candidate(subset, candidate_a, idx_joint_a, kpt_num)
 272                 continue
 273
 274             temp = self._get_connections(candidate_a, candidate_b, score_mid, pafs, threshold)
 275             if not temp:
 276                 continue
 277
 278             if keypoint_id == 0:
 279                 subset = [np.ones(kpt_num) * -1 for _ in temp]
 280                 for i, temp_i in enumerate(temp):
 281                     subset[i][self.limb_seq[0][0] - 1] = temp_i[0]
 282                     subset[i][self.limb_seq[0][1] - 1] = temp_i[1]
 283                     subset[i][-1] = 2
 284                     subset[i][-2] = np.sum(candidates[temp_i[0:2], 2]) + temp_i[2]
 285             else:
 286                 index_a = self.limb_seq[keypoint_id][0] - 1
 287                 index_b = self.limb_seq[keypoint_id][1] - 1
 288                 if keypoint_id in (17, 18):
 289                     subset = self._copy_temperature_to_subset(subset, temp, index_a, index_b)
 290                     continue
 291                 subset = self._add_pose_both_candidates(subset, temp, index_a, index_b, candidates, kpt_num)
 292
 293         return self._filter_subset(subset), candidates
 294
 295     @staticmethod
 296     def get_poses(subset, candidate):
 297         persons_keypoints_x, persons_keypoints_y, persons_keypoints_v = [], [], []
 298         scores = []
 299         for subset_element in subset:
 300             if subset_element.size == 0:
 301                 continue
 302             keypoints_x, keypoints_y, keypoints_v = [0] * 17, [0] * 17, [0] * 17
 303             to_coco_map = [0, -1, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3]
 304             person_score = subset_element[-2]
 305             position_id = -1
 306             for keypoint_id in subset_element[:-2]:
 307                 position_id += 1
 308                 if position_id == 1:  # No 'Neck' in COCO
 309                     continue
 310
 311                 cx, cy, visibility = 0, 0, 0  # Keypoint not found
 312                 if keypoint_id != -1:
 313                     cx, cy = candidate[keypoint_id.astype(int), 0:2]
 314                     cx = cx - 0.5 + 1  # +1 for matlab consistency, coords start from 1
 315                     cy = cy - 0.5 + 1
 316                     visibility = 1
 317                 keypoints_x[to_coco_map[position_id]] = cx
 318                 keypoints_y[to_coco_map[position_id]] = cy
 319                 keypoints_v[to_coco_map[position_id]] = visibility
 320
 321             scores.append(person_score * max(0, (subset_element[-1] - 1)))  # -1 for Neck
 322             persons_keypoints_x.append(keypoints_x)
 323             persons_keypoints_y.append(keypoints_y)
 324             persons_keypoints_v.append(keypoints_v)
 325
 326         persons_keypoints_x = np.array(persons_keypoints_x)
 327         persons_keypoints_y = np.array(persons_keypoints_y)
 328         persons_keypoints_v = np.array(persons_keypoints_v)
 329         scores = np.array(scores)
 330
 331         return persons_keypoints_x, persons_keypoints_y, persons_keypoints_v, scores