def keypoint_results(cls_boxes, pred_heatmaps, ref_boxes): num_classes = cfg.MODEL.NUM_CLASSES cls_keyps = [[] for _ in range(num_classes)] person_idx = keypoint_utils.get_person_class_index() xy_preds = keypoint_utils.heatmaps_to_keypoints(pred_heatmaps, ref_boxes) # NMS OKS if cfg.KRCNN.NMS_OKS: keep = keypoint_utils.nms_oks(xy_preds, ref_boxes, 0.3) xy_preds = xy_preds[keep, :, :] ref_boxes = ref_boxes[keep, :] pred_heatmaps = pred_heatmaps[keep, :, :, :] cls_boxes[person_idx] = cls_boxes[person_idx][keep, :] kps = [xy_preds[i] for i in range(xy_preds.shape[0])] cls_keyps[person_idx] = kps return cls_keyps