smplfitter 0.3.2__py3-none-any.whl

smplfitter/__init__.py

@@ -0,0 +1,21 @@
+"""SMPLFitter provides forward and inverse kinematics for SMPL-family body models.
+
+Main submodules:
+- :mod:`smplfitter.np` - NumPy backend
+- :mod:`smplfitter.pt` - PyTorch backend
+- :mod:`smplfitter.tf` - TensorFlow backend
+- :mod:`smplfitter.jax` - JAX backend
+- :mod:`smplfitter.nb` - Numba backend
+"""
+
+from __future__ import annotations
+
+from .common import ModelData, initialize, _set_module_for_docs
+
+try:
+    from ._version import version as __version__
+except ImportError:
+    __version__ = '0.0.0'
+
+__all__ = ['ModelData', 'initialize', '__version__']
+_set_module_for_docs(__name__, globals(), __all__)

smplfitter/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.3.2'
+__version_tuple__ = version_tuple = (0, 3, 2)
+
+__commit_id__ = commit_id = None

smplfitter/common.py

@@ -0,0 +1,266 @@
+from __future__ import annotations
+
+import contextlib
+import os
+import os.path as osp
+import pickle
+import sys
+import warnings
+from dataclasses import dataclass
+
+import numpy as np
+
+
+def _set_module_for_docs(module_name, module_globals, all_names):
+    """Override __module__ on exported objects so Sphinx resolves package-level names.
+
+    sphinx-codeautolink uses __module__ to find the docs page for a name. Without this,
+    e.g. ``BodyModel`` imported from ``smplfitter.pt.bodymodel`` would not link to
+    ``smplfitter.pt.BodyModel``. The original __module__ is saved as ``_module_original_``
+    so that ``inspect.getsourcefile`` can still find the real source file (see
+    ``module_restored`` in docs/conf.py).
+    """
+    for name in all_names:
+        obj = module_globals.get(name)
+        if obj is not None and callable(obj):
+            obj._module_original_ = obj.__module__
+            obj.__module__ = module_name
+
+
+@dataclass
+class ModelData:
+    """Data loaded from a SMPL-family body model file.
+
+    This dataclass holds all arrays and metadata needed to instantiate a body model
+    in any backend (NumPy, PyTorch, TensorFlow, JAX, Numba).
+    """
+
+    # Tensor data (numpy arrays, to be converted to framework-specific tensors)
+    v_template: np.ndarray
+    """Vertex template in T-pose, shape (num_vertices, 3)."""
+
+    shapedirs: np.ndarray
+    """Shape blend shapes, shape (num_vertices, 3, num_betas)."""
+
+    posedirs: np.ndarray
+    """Pose blend shapes, shape (num_vertices, 3, (num_joints-1)*9)."""
+
+    J_regressor_post_lbs: np.ndarray
+    """Joint regressor for post-LBS joint locations, shape (num_joints, num_vertices)."""
+
+    J_template: np.ndarray
+    """Joint template positions, shape (num_joints, 3)."""
+
+    J_shapedirs: np.ndarray
+    """Joint shape directions, shape (num_joints, 3, num_betas)."""
+
+    kid_shapedir: np.ndarray
+    """Kid shape blend shape for vertices, shape (num_vertices, 3)."""
+
+    kid_J_shapedir: np.ndarray
+    """Kid shape blend shape for joints, shape (num_joints, 3)."""
+
+    weights: np.ndarray
+    """Skinning weights, shape (num_vertices, num_joints)."""
+
+    # Non-tensor data (metadata)
+    kintree_parents: list[int]
+    """Parent joint indices for kinematic tree."""
+
+    faces: np.ndarray
+    """Face indices, shape (num_faces, 3)."""
+
+    num_joints: int
+    """Number of joints in the body model."""
+
+    num_vertices: int
+    """Number of vertices in the body model mesh."""
+
+    vertex_subset: np.ndarray
+    """Indices of vertices used (for partial models)."""
+
+
+def initialize(
+    model_name,
+    gender,
+    model_root=None,
+    num_betas=None,
+    vertex_subset_size=None,
+    vertex_subset=None,
+    faces=None,
+    joint_regressor_post_lbs=None,
+):
+    if model_root is None:
+        body_models_dir = os.getenv('SMPLFITTER_BODY_MODELS')
+        if body_models_dir is None:
+            data_root = os.getenv('DATA_ROOT', '.')
+            body_models_dir = f'{data_root}/body_models'
+        model_root = f'{body_models_dir}/{model_name}'
+
+    with monkey_patched_for_chumpy():
+        if model_name == 'smpl':
+            gender_str = dict(f='f', m='m', n='neutral')[gender[0]]
+            filename = f'basicmodel_{gender_str}_lbs_10_207_0_v1.1.0.pkl'
+        elif model_name in ('smplx', 'smplxlh', 'smplxmoyo'):
+            gender_str = dict(f='FEMALE', m='MALE', n='NEUTRAL')[gender[0]]
+            filename = f'SMPLX_{gender_str}.npz'
+        elif model_name == 'smplh':
+            gender_str = dict(f='female', m='male')[gender[0]]
+            filename = f'SMPLH_{gender_str}.pkl'
+        elif model_name == 'smplh16':
+            gender_str = dict(f='female', m='male', n='neutral')[gender[0]]
+            filename = osp.join(gender_str, 'model.npz')
+        elif model_name == 'mano':
+            filename = 'MANO_RIGHT.pkl'
+        else:
+            raise ValueError(f'Unknown model name: {model_name}')
+
+        filepath = osp.join(model_root, filename)
+        try:
+            if filename.endswith('.npz'):
+                smpl_data = np.load(filepath)
+            else:
+                with open(filepath, 'rb') as f:
+                    smpl_data = pickle.load(f, encoding='latin1')
+        except FileNotFoundError:
+            raise FileNotFoundError(
+                f'Body model file not found: {filepath}\n\n'
+                f'Set the body model location using one of:\n'
+                f"  1. BodyModel('{model_name}', '{gender}', "
+                f"model_root='/your/path/body_models/{model_name}')\n"
+                f'  2. export SMPLFITTER_BODY_MODELS=/your/path/body_models\n'
+                f'  3. export DATA_ROOT=/your/path   '
+                f'(looks for $DATA_ROOT/body_models/)\n\n'
+                f'Download models: python -m smplfitter.download\n'
+                f'Register first at https://smpl.is.tue.mpg.de/'
+            ) from None
+
+    res = {}
+    res['shapedirs'] = np.array(smpl_data['shapedirs'], dtype=np.float64)
+    res['posedirs'] = np.array(smpl_data['posedirs'], dtype=np.float64)
+    res['v_template'] = np.array(smpl_data['v_template'], dtype=np.float64)
+
+    if not isinstance(smpl_data['J_regressor'], np.ndarray):
+        res['J_regressor'] = np.array(smpl_data['J_regressor'].toarray(), dtype=np.float64)
+    else:
+        res['J_regressor'] = smpl_data['J_regressor'].astype(np.float64)
+
+    res['weights'] = np.array(smpl_data['weights'])
+    res['faces'] = np.array(smpl_data['f'].astype(np.int32))
+    res['kintree_parents'] = smpl_data['kintree_table'][0].tolist()
+    res['num_joints'] = len(res['kintree_parents'])
+    res['num_vertices'] = len(res['v_template'])
+
+    # Kid model has an additional shape parameter which pulls the mesh towards the SMIL mean
+    # template
+    if model_name.lower().startswith('smpl'):
+        kid_path = os.path.join(model_root, 'kid_template.npy')
+        try:
+            v_template_smil = np.load(kid_path).astype(np.float64)
+        except FileNotFoundError:
+            raise FileNotFoundError(
+                f'Kid template not found: {kid_path}\n'
+                f'Download it: python -m smplfitter.download'
+            ) from None
+        res['kid_shapedir'] = (
+            v_template_smil - np.mean(v_template_smil, axis=0) - res['v_template']
+        )
+        res['kid_J_shapedir'] = res['J_regressor'] @ res['kid_shapedir']
+    else:
+        res['kid_shapedir'] = np.zeros_like(res['v_template'])
+        res['kid_J_shapedir'] = np.zeros((res['num_joints'], 3))
+
+    if 'J_shapedirs' in smpl_data:
+        res['J_shapedirs'] = np.array(smpl_data['J_shapedirs'], dtype=np.float64)
+    else:
+        res['J_shapedirs'] = np.einsum('jv,vcs->jcs', res['J_regressor'], res['shapedirs'])
+
+    if 'J_template' in smpl_data:
+        res['J_template'] = np.array(smpl_data['J_template'], dtype=np.float64)
+    else:
+        res['J_template'] = res['J_regressor'] @ res['v_template']
+
+    res['v_template'] = res['v_template'] - np.einsum(
+        'vcx,x->vc',
+        res['posedirs'],
+        np.reshape(np.tile(np.eye(3, dtype=np.float64), [res['num_joints'] - 1, 1]), [-1]),
+    )
+
+    if vertex_subset_size is not None:
+        vertex_subset_dict = np.load(f'{model_root}/vertex_subset_{vertex_subset_size}.npz')
+        vertex_subset = vertex_subset_dict['i_verts']
+        faces = vertex_subset_dict['faces']
+        joint_regressor_post_lbs = np.load(
+            f'{model_root}/vertex_subset_joint_regr_post_lbs_{vertex_subset_size}.npy'
+        )
+
+    if vertex_subset is None:
+        vertex_subset = np.arange(res['num_vertices'], dtype=np.int64)
+    else:
+        vertex_subset = np.array(vertex_subset, dtype=np.int64)
+
+    if faces is None:
+        faces = res['faces']
+
+    if joint_regressor_post_lbs is None:
+        joint_regressor_post_lbs = res['J_regressor']
+
+    return ModelData(
+        v_template=res['v_template'][vertex_subset],
+        shapedirs=res['shapedirs'][vertex_subset, :, :num_betas],
+        posedirs=res['posedirs'][vertex_subset],
+        J_regressor_post_lbs=joint_regressor_post_lbs,
+        J_template=res['J_template'],
+        J_shapedirs=res['J_shapedirs'][:, :, :num_betas],
+        kid_shapedir=res['kid_shapedir'][vertex_subset],
+        kid_J_shapedir=res['kid_J_shapedir'],
+        weights=res['weights'][vertex_subset],
+        kintree_parents=res['kintree_parents'],
+        faces=faces,
+        num_joints=res['num_joints'],
+        num_vertices=len(vertex_subset),
+        vertex_subset=vertex_subset,
+    )
+
+
+@contextlib.contextmanager
+def monkey_patched_for_chumpy():
+    """The pickle file of SMPLH imports chumpy and it tries to import np.bool etc which are
+    not available anymore.
+    """
+    added = []
+    for name in ['bool', 'int', 'object', 'str']:
+        if name not in dir(np):
+            try:
+                sys.modules[f'numpy.{name}'] = getattr(np, name + '_')
+                added.append(name)
+            except AttributeError:
+                pass
+
+    sys.modules['numpy.float'] = float
+    sys.modules['numpy.complex'] = np.complex128
+    sys.modules['numpy.NINF'] = -np.inf
+    np.NINF = -np.inf  # type: ignore[misc]
+    np.complex = np.complex128  # type: ignore[misc]
+    np.float = float  # type: ignore[misc]
+
+    if 'unicode' not in dir(np):
+        sys.modules['numpy.unicode'] = np.str_
+        added.append('unicode')
+
+    import inspect
+
+    added_getargspec = False
+    if not hasattr(inspect, 'getargspec'):
+        inspect.getargspec = inspect.getfullargspec
+        added_getargspec = True
+
+    with warnings.catch_warnings():
+        warnings.simplefilter('ignore', FutureWarning)
+        yield
+
+    for name in added:
+        del sys.modules[f'numpy.{name}']
+
+    if added_getargspec:
+        del inspect.getargspec

smplfitter/decimation/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

smplfitter/decimation/decimate_body_models.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+import os
+
+import numpy as np
+from .. import np as smplfitter_np
+import trimesh
+from scipy.optimize import linear_sum_assignment
+from scipy.spatial.distance import cdist
+
+
+def main():
+    DATA_ROOT = os.getenv('DATA_ROOT', default='.')
+    for model_name in ['smpl', 'smplx']:
+        model_root = f'{DATA_ROOT}/body_models/{model_name}'
+        bm = smplfitter_np.get_cached_body_model(model_name)
+        verts = bm.single()['vertices']
+
+        for n in [32, 64, 128, 256, 512, 1024]:
+            i_verts, faces = decimate(verts, bm.faces, n)
+            if len(i_verts) != n:
+                print(f'Failed to decimate to {n} vertices')
+                continue
+            else:
+                print(f'Decimated to {n} vertices')
+            np.savez(f'{model_root}/vertex_subset_{n}.npz', i_verts=i_verts, faces=faces)
+
+
+def _decimate(verts, faces, n_verts_out=128):
+    n_faces = 2 * n_verts_out - 4
+    decimated_mesh = trimesh.Trimesh(verts, faces).simplify_quadric_decimation(face_count=n_faces)
+    row_ind, col_ind = linear_sum_assignment(cdist(verts, decimated_mesh.vertices))
+    i_verts = row_ind[np.argsort(col_ind)]
+    return i_verts, decimated_mesh.faces
+
+
+def decimate(verts, faces, n_verts_out=128, n_trials=100):
+    # We may have to try it multiple times to get the desired number of vertices
+    # since the decimation algorithm may not always return the exact number of vertices we ask.
+    # The documentation explicitly states this.
+
+    n_verts_out_arg = n_verts_out
+    i_verts, new_faces = _decimate(verts, faces, n_verts_out_arg)
+
+    for _ in range(n_trials):
+        if i_verts.shape[0] == n_verts_out:
+            break
+
+        n_verts_out_arg += n_verts_out - i_verts.shape[0]
+        i_verts, new_faces = _decimate(verts, faces, n_verts_out_arg)
+
+    return i_verts, new_faces
+
+
+if __name__ == '__main__':
+    main()

smplfitter/decimation/make_post_lbs_joint_regressors.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+# This module finds joint regressors for subsets of body model vertices.
+# That is, it finds a linear regressor that maps a subset of vertices to the joints.
+# And this is done for the post-LBS case, i.e. the vertices and joints are posed with linear
+# blend skinning.
+# This is useful in body model fitting in case the user does not provide the joint positions.
+# In that case we must run a regressor on the (already posed) vertices.
+#
+# The joint regressors need to create convex combinations, i.e. the weights must be non-negative
+# and sum to 1 for each joint. They should also be sparse and spatially compact.
+# That is, each joint should only depend on a few vertices that are spatially close to each other.
+# This is achieved by adding a regularization term that encourages the weights to be
+# sparse (L1/2 norm), and a term that computes the weighted variance of the template vertices.
+
+import os
+
+import numpy as np
+from .. import pt as smplfitter_pt
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, IterableDataset
+from tqdm import tqdm
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+DATA_ROOT = os.getenv('DATA_ROOT', default='.')
+torch.backends.cudnn.benchmark = True
+
+
+def main():
+    for body_model_name in ['smpl', 'smplx']:
+        canonical_verts = torch.from_numpy(
+            np.load(f'{DATA_ROOT}/nlf/canonical_vertices_{body_model_name}.npy')
+        ).float()
+        body_model = smplfitter_pt.BodyModel(body_model_name).to(DEVICE)
+        dataset = RandomBodyParamDataset(num_betas=16, num_joints=body_model.num_joints)
+        dataloader = DataLoader(dataset, batch_size=None, num_workers=4, pin_memory=True)
+
+        for n_verts_subset in reversed(
+            [32, 64, 128, 256, 512, 1024, 2048, 4096, body_model.num_vertices]
+        ):
+            print(f'Fitting joint regressor for {n_verts_subset} vertices')
+            out_path = (
+                f'{DATA_ROOT}/body_models/{body_model_name}/'
+                f'vertex_subset_joint_regr_post_lbs_{n_verts_subset}.npy'
+            )
+            # if os.path.exists(out_path):
+            #     continue
+
+            if n_verts_subset == body_model.num_vertices:
+                i_verts = torch.arange(body_model.num_vertices)
+            else:
+                subset = np.load(
+                    f'{DATA_ROOT}/body_models/{body_model_name}/vertex_subset_{n_verts_subset}.npz'
+                )
+                i_verts = torch.from_numpy(subset['i_verts'])
+
+            model = ConvexCombiningRegressor(len(i_verts), body_model.num_joints).to(DEVICE)
+            trainer = ConvexCombiningRegressorTrainer(
+                model=model,
+                body_model=body_model,
+                template_verts=canonical_verts[i_verts],
+                vertex_subset=i_verts,
+                regul_lambda=6e-5,
+            )
+            optimizer = optim.Adam(model.parameters(), lr=1e0)
+            scheduler = optim.lr_scheduler.LambdaLR(
+                optimizer, lr_lambda=lambda step: 1.0 if step < int(37500 * 0.9) else 1e-3
+            )
+            trainer.train_loop(
+                dataloader, total_steps=37500, optimizer=optimizer, scheduler=scheduler
+            )
+            model.threshold_for_sparsity(1e-3)
+
+            scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: 30)
+            trainer.train_loop(
+                dataloader, total_steps=37500 + 12500, optimizer=optimizer, scheduler=scheduler
+            )
+            model.threshold_for_sparsity(1e-3)
+
+            J_subset = model.get_w().cpu().detach().numpy().T
+            print(f'Sparsity ratio: {sparsity_ratio(J_subset)}')
+            np.save(out_path, J_subset)
+
+
+class ConvexCombiningRegressor(nn.Module):
+    def __init__(self, n_in_points, n_out_points):
+        super().__init__()
+        self.n_in = n_in_points
+        self.n_out = n_out_points
+
+        self.weight_mask = nn.Parameter(torch.ones(n_in_points, n_out_points), requires_grad=False)
+        self.weights = nn.Parameter(torch.empty(n_in_points, n_out_points))
+        nn.init.uniform_(self.weights, -1, 1)
+
+    def forward(self, x):
+        w = self.get_w()
+        return torch.einsum('bjc,jJ->bJc', x, w)
+
+    def get_w(self):
+        w = torch.nn.functional.softplus(self.weights) * self.weight_mask
+        return w / w.sum(dim=0, keepdim=True)
+
+    def threshold_for_sparsity(self, threshold=1e-3):
+        with torch.no_grad():
+            self.weight_mask.data = (torch.abs(self.get_w()) > threshold).float()
+
+
+class ConvexCombiningRegressorTrainer:
+    def __init__(self, model, body_model, template_verts, vertex_subset, regul_lambda=3e-5):
+        self.model = model
+        self.body_model = body_model
+        self.template_verts = template_verts.to(DEVICE)
+        self.vertex_subset = vertex_subset
+        self.regul_lambda = regul_lambda
+        self.current_step = 0
+
+    def train_loop(self, dataloader, total_steps, optimizer, scheduler=None):
+        progress_bar = tqdm(
+            total=total_steps,
+            desc='Training',
+            unit='step',
+            dynamic_ncols=True,
+            initial=self.current_step,
+        )
+
+        for batch in dataloader:
+            batch = EasyDict({k: v.to(DEVICE) for k, v in batch.items()})
+            need_metrics = self.current_step % 100 == 0
+            losses, metrics = self.train_step(
+                batch, optimizer, scheduler, need_metrics=need_metrics
+            )
+            self.current_step += 1
+            progress_bar.update(1)
+
+            if need_metrics:
+                losses_and_metrics_str = ', '.join(
+                    f'{k}: {v.item():.4f}' for k, v in (losses | metrics).items()
+                )
+                progress_bar.set_postfix_str(losses_and_metrics_str)
+
+            if self.current_step >= total_steps:
+                break
+
+    def forward_train(self, inps):
+        preds = EasyDict()
+        r = self.body_model(inps.pose, inps.shape)
+        preds.postjoints = self.body_model.J_regressor_post_lbs @ r['vertices']
+        preds.pose3d = self.model(r['vertices'][:, self.vertex_subset])
+        preds.pose3d_gt = r['joints']
+        return preds
+
+    def train_step(self, inps, optimizer, scheduler=None, need_metrics=True):
+        self.model.train()
+        optimizer.zero_grad()
+        preds = self.forward_train(inps)
+        losses = self.compute_losses(inps, preds)
+        losses.loss.backward()
+        optimizer.step()
+        if scheduler is not None:
+            scheduler.step()
+
+        if need_metrics:
+            with torch.no_grad():
+                metrics = self.compute_metrics(inps, preds)
+        else:
+            metrics = None
+
+        return losses, metrics
+
+    def compute_losses(self, inps, preds):
+        losses = EasyDict()
+        losses.main_loss = torch.mean(torch.abs(preds.pose3d_gt - preds.pose3d))
+
+        w = self.model.get_w()
+        losses.regul = torch.sum(soft_sqrt(torch.abs(w), 1e-5)) / w.shape[1]
+        # losses.supp = mean_spatial_support(self.template_verts, w)
+        losses.loss = (
+            losses.main_loss + self.regul_lambda * losses.regul  # + self.supp_lambda * losses.supp
+        )
+        return losses
+
+    def compute_metrics(self, inps, preds):
+        m = EasyDict()
+        dist = torch.norm(preds.pose3d_gt - preds.pose3d, dim=-1)
+        m.pck1 = pck(dist, 0.01)
+        m.pck2 = pck(dist, 0.02)
+        m.pck3 = pck(dist, 0.03)
+        m.pck7 = pck(dist, 0.07)
+        m.euclidean = torch.mean(dist)
+        m.l1 = torch.mean(torch.abs(preds.pose3d_gt - preds.pose3d))
+        m.max_supp = torch.max(
+            torch.sqrt(spatial_support(self.template_verts, self.model.get_w()))
+        )
+
+        dist = torch.norm(preds.pose3d_gt - preds.postjoints, dim=-1)
+        m.pck1_post = pck(dist, 0.01)
+        return m
+
+
+class RandomBodyParamDataset(IterableDataset):
+    def __init__(self, batch_size=72, num_joints=24, num_betas=16):
+        self.batch_size = batch_size
+        self.num_joints = num_joints
+        self.num_betas = num_betas
+        self.rng = torch.Generator()
+
+    def __iter__(self):
+        with torch.no_grad():
+            while True:
+                pose = self.random_smpl_pose() / 3
+                shape = torch.empty((self.batch_size, self.num_betas))
+                nn.init.trunc_normal_(shape, mean=0.0, std=1.0, a=-2.0, b=2.0, generator=self.rng)
+                yield {'pose': pose, 'shape': shape * 10}
+
+    def random_rotvec(self, batch_size):
+        rand = torch.rand(3, batch_size, generator=self.rng)
+        r1 = torch.sqrt(1 - rand[0])
+        r2 = torch.sqrt(rand[0])
+        t1 = 2 * np.pi * rand[1]
+        t2 = 2 * np.pi * rand[2]
+        cost2 = torch.cos(t2)
+        xyz = torch.stack([torch.sin(t1) * r1, torch.cos(t1) * r1, torch.sin(t2) * r2], dim=0)
+        return (xyz / torch.sqrt(1 - cost2**2 * rand[0]) * 2 * torch.acos(cost2 * r2)).T
+
+    def random_smpl_pose(self):
+        return self.random_rotvec(self.batch_size * self.num_joints).reshape(
+            self.batch_size, self.num_joints * 3
+        )
+
+
+def sparsity_ratio(J):
+    return np.count_nonzero(np.abs(J) > 1e-4) / J.shape[0]
+
+
+def soft_sqrt(x, eps):
+    return x / torch.sqrt(x + eps)
+
+
+def spatial_support(template, weights):
+    weighted_mean = torch.matmul(weights.t(), template)
+    sq_diff = (template[np.newaxis, :] - weighted_mean[:, np.newaxis]) ** 2
+    sq_dists = torch.sum(sq_diff, dim=-1)
+    return torch.einsum('Jj,jJ->J', sq_dists, torch.abs(weights))
+
+
+def pck(x, t):
+    return (x <= t).float().mean()
+
+
+class EasyDict(dict):
+    def __init__(self, d=None, **kwargs):
+        super().__init__()
+        if d is None:
+            d = {}
+        else:
+            d = dict(d)
+        if kwargs:
+            d.update(**kwargs)
+        for k, v in d.items():
+            setattr(self, k, v)
+        # Class attributes
+        for k in self.__class__.__dict__.keys():
+            if not (k.startswith('__') and k.endswith('__')) and k not in ('update', 'pop'):
+                setattr(self, k, getattr(self, k))
+
+    def __setattr__(self, name, value):
+        super().__setitem__(name, value)
+
+    def __getattr__(self, name):
+        try:
+            return super().__getitem__(name)
+        except KeyError:
+            raise AttributeError(name)
+
+    __setitem__ = __setattr__
+
+    def update(self, e=None, **f):
+        d = e or dict()
+        d.update(f)
+        for k, v in d.items():
+            setattr(self, k, v)
+
+    def pop(self, k, d=None):
+        delattr(self, k)
+        return super().pop(k, d)
+
+
+if __name__ == '__main__':
+    main()