Compute the losses related to the keypoints.
(self, outputs, targets, indices,
idx, num_boxes, data_batch,
face_hand_kpt=False)
| 130 | return losses |
| 131 | |
| 132 | def loss_keypoints(self, outputs, targets, indices, |
| 133 | idx, num_boxes, data_batch, |
| 134 | face_hand_kpt=False): |
| 135 | """Compute the losses related to the keypoints.""" |
| 136 | indices = indices[0] |
| 137 | device = outputs['pred_logits'].device |
| 138 | losses = {} |
| 139 | |
| 140 | src_body_keypoints = outputs['pred_keypoints'][idx] # xyxyvv |
| 141 | if len(src_body_keypoints) == 0: |
| 142 | losses.append({ |
| 143 | 'loss_keypoints': src_body_keypoints.sum() * 0 + \ |
| 144 | outputs['pred_smpl_cam'][idx].float().sum()*0, |
| 145 | 'loss_oks': src_body_keypoints.sum() * 0, |
| 146 | }) |
| 147 | else: |
| 148 | Z_pred = src_body_keypoints[:, 0:(self.num_body_points * 2)] # [2, 2*14] |
| 149 | V_pred = src_body_keypoints[:, (self.num_body_points * 2):] |
| 150 | targets_body_keypoints = torch.cat( |
| 151 | [t['keypoints'][i] for t, (_, i) in zip(targets, indices)], |
| 152 | dim=0) |
| 153 | targets_area = torch.cat( |
| 154 | [t['area'][i] for t, (_, i) in zip(targets, indices)], dim=0) |
| 155 | target_body_boxes_conf = torch.cat( |
| 156 | [t[i] for t, (_, i) in zip(data_batch['body_bbox_valid'], indices)], dim=0) |
| 157 | Z_gt = targets_body_keypoints[:, 0:(self.num_body_points * 2)] |
| 158 | V_gt: torch.Tensor = targets_body_keypoints[:, (self.num_body_points * 2):] |
| 159 | body_kps_conf = V_gt.sum(-1)>0 |
| 160 | body_num_boxes = (body_kps_conf * target_body_boxes_conf).sum() |
| 161 | oks_loss = self.body_oks(Z_pred, |
| 162 | Z_gt, |
| 163 | V_gt, |
| 164 | targets_area, |
| 165 | weight=None, |
| 166 | avg_factor=None, |
| 167 | reduction_override=None) |
| 168 | oks_loss*= body_kps_conf * target_body_boxes_conf |
| 169 | pose_loss = F.l1_loss(Z_pred, Z_gt, reduction='none') |
| 170 | pose_loss = pose_loss * V_gt.repeat_interleave(2, dim=1) |
| 171 | pose_loss = pose_loss.sum(-1) * target_body_boxes_conf |
| 172 | if body_num_boxes>0: |
| 173 | losses['loss_keypoints'] = pose_loss.sum() / body_num_boxes |
| 174 | losses['loss_oks'] = oks_loss.sum() / body_num_boxes |
| 175 | else: |
| 176 | losses['loss_keypoints'] = src_body_keypoints.sum() * 0 |
| 177 | losses['loss_oks'] = src_body_keypoints.sum() * 0 |
| 178 | |
| 179 | # lhand |
| 180 | if 'pred_lhand_keypoints' in outputs and face_hand_kpt: |
| 181 | src_lhand_keypoints = outputs['pred_lhand_keypoints'][idx] # xyxyvv |
| 182 | if len(src_lhand_keypoints) == 0: |
| 183 | losses.update({ |
| 184 | 'loss_lhand_keypoints': src_lhand_keypoints.sum() * 0, |
| 185 | 'loss_lhand_oks':src_lhand_keypoints.sum() * 0, |
| 186 | }) |
| 187 | else: |
| 188 | Z_pred = src_lhand_keypoints[:, 0:(self.num_hand_points * 2)] # [2, 2*14] |
| 189 | V_pred = src_lhand_keypoints[:, (self.num_hand_points * 2):] |