98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
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]))
|