Flip pose heatmap according to joint pairs. Parameters ---------- heatmap : numpy.ndarray Heatmap of joints. joint_pairs : list List of joint pairs. shift : bool Whether to shift the output. Returns ------- numpy.ndarray Flipped heatm
(heatmap, joint_pairs, shift=False)
| 514 | |
| 515 | |
| 516 | def flip_heatmap(heatmap, joint_pairs, shift=False): |
| 517 | """Flip pose heatmap according to joint pairs. |
| 518 | |
| 519 | Parameters |
| 520 | ---------- |
| 521 | heatmap : numpy.ndarray |
| 522 | Heatmap of joints. |
| 523 | joint_pairs : list |
| 524 | List of joint pairs. |
| 525 | shift : bool |
| 526 | Whether to shift the output. |
| 527 | |
| 528 | Returns |
| 529 | ------- |
| 530 | numpy.ndarray |
| 531 | Flipped heatmap. |
| 532 | |
| 533 | """ |
| 534 | assert (heatmap.dim() == 3 or heatmap.dim() == 4) |
| 535 | out = flip(heatmap) |
| 536 | |
| 537 | for pair in joint_pairs: |
| 538 | dim0, dim1 = pair |
| 539 | idx = torch.Tensor((dim0, dim1)).long() |
| 540 | inv_idx = torch.Tensor((dim1, dim0)).long() |
| 541 | if out.dim() == 4: |
| 542 | out[:, idx] = out[:, inv_idx] |
| 543 | else: |
| 544 | out[idx] = out[inv_idx] |
| 545 | |
| 546 | if shift: |
| 547 | if out.dim() == 3: |
| 548 | out[:, :, 1:] = out[:, :, 0:-1] |
| 549 | else: |
| 550 | out[:, :, :, 1:] = out[:, :, :, 0:-1] |
| 551 | return out |
| 552 | |
| 553 | |
| 554 | def flip_coord(preds, joint_pairs, width_dim, shift=False, flatten=True): |