Spaces:
Build error
Build error
File size: 3,193 Bytes
4409449 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
from typing import List
import torch
from torch import Tensor
from torchmetrics import Metric
from .utils import *
# motion reconstruction metric
class MRMetrics(Metric):
def __init__(self,
njoints,
jointstype: str = "mmm",
force_in_meter: bool = True,
align_root: bool = True,
dist_sync_on_step=True,
**kwargs):
super().__init__(dist_sync_on_step=dist_sync_on_step)
self.name = 'Motion Reconstructions'
self.jointstype = jointstype
self.align_root = align_root
self.force_in_meter = force_in_meter
self.add_state("count", default=torch.tensor(0), dist_reduce_fx="sum")
self.add_state("count_seq",
default=torch.tensor(0),
dist_reduce_fx="sum")
self.add_state("MPJPE",
default=torch.tensor([0.0]),
dist_reduce_fx="sum")
self.add_state("PAMPJPE",
default=torch.tensor([0.0]),
dist_reduce_fx="sum")
self.add_state("ACCEL",
default=torch.tensor([0.0]),
dist_reduce_fx="sum")
# todo
# self.add_state("ROOT", default=torch.tensor([0.0]), dist_reduce_fx="sum")
self.MR_metrics = ["MPJPE", "PAMPJPE", "ACCEL"]
# All metric
self.metrics = self.MR_metrics
def compute(self, sanity_flag):
if self.force_in_meter:
# different jointstypes have different scale factors
# if self.jointstype == 'mmm':
# factor = 1000.0
# elif self.jointstype == 'humanml3d':
# factor = 1000.0 * 0.75 / 480
factor = 1000.0
else:
factor = 1.0
count = self.count
count_seq = self.count_seq
mr_metrics = {}
mr_metrics["MPJPE"] = self.MPJPE / count * factor
mr_metrics["PAMPJPE"] = self.PAMPJPE / count * factor
# accel error: joints_gt[:-2] - 2 * joints_gt[1:-1] + joints_gt[2:]
# n-2 for each sequences
mr_metrics["ACCEL"] = self.ACCEL / (count - 2 * count_seq) * factor
# Reset
self.reset()
return mr_metrics
def update(self, joints_rst: Tensor, joints_ref: Tensor,
lengths: List[int]):
assert joints_rst.shape == joints_ref.shape
assert joints_rst.dim() == 4
# (bs, seq, njoint=22, 3)
self.count += sum(lengths)
self.count_seq += len(lengths)
# avoid cuda error of DDP in pampjpe
rst = joints_rst.detach().cpu()
ref = joints_ref.detach().cpu()
# align root joints index
if self.align_root and self.jointstype in ['mmm', 'humanml3d']:
align_inds = [0]
else:
align_inds = None
for i in range(len(lengths)):
self.MPJPE += torch.sum(
calc_mpjpe(rst[i], ref[i], align_inds=align_inds))
self.PAMPJPE += torch.sum(calc_pampjpe(rst[i], ref[i]))
self.ACCEL += torch.sum(calc_accel(rst[i], ref[i]))
|