Spaces:
Sleeping
Sleeping
| 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])) | |