src/smal_pytorch/smal_model/smal_torch_new.py

"""
PyTorch implementation of the SMAL/SMPL model
see:
    1.) https://github.com/silviazuffi/smalst/blob/master/smal_model/smal_torch.py
    2.) https://github.com/benjiebob/SMALify/blob/master/smal_model/smal_torch.py
main changes compared to SMALST and WLDO:
    * new model
        (/ps/scratch/nrueegg/new_projects/side_packages/SMALify/new_smal_pca/results/my_tposeref_results_3/)
        dogs are part of the pca to create the model
        al meshes are centered around their root joint
        the animals are all scaled such that their body length (butt to breast) is 1
            X_init = np.concatenate((vertices_dogs, vertices_smal), axis=0)     # vertices_dogs
            X = []
            for ind in range(0, X_init.shape[0]):
                X_tmp, _, _, _ = align_smal_template_to_symmetry_axis(X_init[ind, :, :], subtract_mean=True)   # not sure if this is necessary
                X.append(X_tmp)
            X = np.asarray(X)
            # define points which will be used for normalization
            idxs_front = [6, 16, 8, 964]      #  [1172, 6, 16, 8, 964]
            idxs_back = [174, 2148, 175, 2149]       # not in the middle, but pairs
            reg_j = np.asarray(dd['J_regressor'].todense())
            # normalize the meshes such that X_frontback_dist is 1 and the root joint is in the center (0, 0, 0)
            X_front = X[:, idxs_front, :].mean(axis=1)
            X_back = X[:, idxs_back, :].mean(axis=1)
            X_frontback_dist = np.sqrt(((X_front - X_back)**2).sum(axis=1))
            X = X / X_frontback_dist[:, None, None]
            X_j0 = np.sum(X[:, reg_j[0, :]>0, :] * reg_j[0, (reg_j[0, :]>0)][None, :, None], axis=1)
            X = X - X_j0[:, None, :]
    * add limb length changes the same way as in WLDO
    * overall scale factor is added
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import torch
import chumpy as ch
import os.path
from torch import nn
from torch.autograd import Variable
import pickle as pkl 
from .batch_lbs import batch_rodrigues, batch_global_rigid_transformation, batch_global_rigid_transformation_biggs, get_bone_length_scales, get_beta_scale_mask

from .smal_basics import align_smal_template_to_symmetry_axis, get_symmetry_indices

import os
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from configs.SMAL_configs import KEY_VIDS, CANONICAL_MODEL_JOINTS, IDXS_BONES_NO_REDUNDANCY, SMAL_MODEL_PATH
 
from smal_pytorch.utils import load_vertex_colors


# There are chumpy variables so convert them to numpy.
def undo_chumpy(x):
    return x if isinstance(x, np.ndarray) else x.r

# class SMAL(object):
class SMAL(nn.Module):        
    def __init__(self, pkl_path=SMAL_MODEL_PATH, n_betas=None, template_name='neutral', use_smal_betas=True, logscale_part_list=None):
        super(SMAL, self).__init__()

        if logscale_part_list is None:
            self.logscale_part_list = ['legs_l', 'legs_f', 'tail_l', 'tail_f', 'ears_y', 'ears_l', 'head_l'] 
        self.betas_scale_mask = get_beta_scale_mask(part_list=self.logscale_part_list)
        self.num_betas_logscale = len(self.logscale_part_list)

        self.use_smal_betas = use_smal_betas

        # -- Load SMPL params --
        try:
            with open(pkl_path, 'r') as f:
                dd = pkl.load(f)
        except (UnicodeDecodeError, TypeError) as e:
            with open(pkl_path, 'rb') as file:
                u = pkl._Unpickler(file)
                u.encoding = 'latin1'
                dd = u.load()

        self.f = dd['f']
        self.register_buffer('faces', torch.from_numpy(self.f.astype(int)))

        # get the correct template (mean shape)
        if template_name=='neutral':
            v_template = dd['v_template'] 
            v = v_template
        else:
            raise NotImplementedError

        # Mean template vertices
        self.register_buffer('v_template', torch.Tensor(v))
        # Size of mesh [Number of vertices, 3]
        self.size = [self.v_template.shape[0], 3]
        self.num_betas = dd['shapedirs'].shape[-1]
        # symmetry indices
        self.sym_ids_dict = get_symmetry_indices()

        # Shape blend shape basis
        shapedir = np.reshape(undo_chumpy(dd['shapedirs']), [-1, self.num_betas]).T
        shapedir.flags['WRITEABLE'] = True      # not sure why this is necessary
        self.register_buffer('shapedirs', torch.Tensor(shapedir))

        # Regressor for joint locations given shape 
        self.register_buffer('J_regressor', torch.Tensor(dd['J_regressor'].T.todense()))

        # Pose blend shape basis
        num_pose_basis = dd['posedirs'].shape[-1]
        
        posedirs = np.reshape(undo_chumpy(dd['posedirs']), [-1, num_pose_basis]).T
        self.register_buffer('posedirs', torch.Tensor(posedirs))
        
        # indices of parents for each joints
        self.parents = dd['kintree_table'][0].astype(np.int32)

        # LBS weights
        self.register_buffer('weights', torch.Tensor(undo_chumpy(dd['weights'])))


    def _caclulate_bone_lengths_from_J(self, J, betas_logscale):
        # NEW: calculate bone lengths:
        all_bone_lengths_list = []
        for i in range(1, self.parents.shape[0]):
            bone_vec = J[:, i] - J[:, self.parents[i]]
            bone_length = torch.sqrt(torch.sum(bone_vec ** 2, axis=1))
            all_bone_lengths_list.append(bone_length)
        all_bone_lengths = torch.stack(all_bone_lengths_list)
        # some bones are pairs, it is enough to take one of the two bones  
        all_bone_length_scales = get_bone_length_scales(self.logscale_part_list, betas_logscale)
        all_bone_lengths = all_bone_lengths.permute((1,0)) * all_bone_length_scales

        return all_bone_lengths     #.permute((1,0))      


    def caclulate_bone_lengths(self, beta, betas_logscale, shapedirs_sel=None, short=True):
        nBetas = beta.shape[1]

        # 1. Add shape blend shapes
        # do we use the original shapedirs or a new set of selected shapedirs?
        if shapedirs_sel is None:
            shapedirs_sel = self.shapedirs[:nBetas,:]
        else: 
            assert shapedirs_sel.shape[0] == nBetas
        v_shaped = self.v_template + torch.reshape(torch.matmul(beta, shapedirs_sel), [-1, self.size[0], self.size[1]])

        # 2. Infer shape-dependent joint locations.
        Jx = torch.matmul(v_shaped[:, :, 0], self.J_regressor)
        Jy = torch.matmul(v_shaped[:, :, 1], self.J_regressor)
        Jz = torch.matmul(v_shaped[:, :, 2], self.J_regressor)
        J = torch.stack([Jx, Jy, Jz], dim=2)

        # calculate bone lengths
        all_bone_lengths = self._caclulate_bone_lengths_from_J(J, betas_logscale)
        selected_bone_lengths = all_bone_lengths[:, IDXS_BONES_NO_REDUNDANCY]
        
        if short:
            return selected_bone_lengths
        else:
            return all_bone_lengths



    def __call__(self, beta, betas_limbs, theta=None, pose=None, trans=None, del_v=None, get_skin=True, keyp_conf='red', get_all_info=False, shapedirs_sel=None):
        device = beta.device

        betas_logscale = betas_limbs
        # NEW: allow that rotation is given as rotation matrices instead of axis angle rotation
        #   theta: BSxNJointsx3 or BSx(NJoints*3)
        #   pose: NxNJointsx3x3
        if (theta is None) and (pose is None):
            raise ValueError("Either pose (rotation matrices NxNJointsx3x3) or theta (axis angle BSxNJointsx3) must be given")
        elif (theta is not None) and (pose is not None):
            raise ValueError("Not both pose (rotation matrices NxNJointsx3x3) and theta (axis angle BSxNJointsx3) can be given")

        if True:        # self.use_smal_betas:
            nBetas = beta.shape[1]
        else:
            nBetas = 0

        # 1. Add shape blend shapes
        # do we use the original shapedirs or a new set of selected shapedirs?
        if shapedirs_sel is None:
            shapedirs_sel = self.shapedirs[:nBetas,:]
        else: 
            assert shapedirs_sel.shape[0] == nBetas
        
        if nBetas > 0:
            if del_v is None:
                v_shaped = self.v_template + torch.reshape(torch.matmul(beta, shapedirs_sel), [-1, self.size[0], self.size[1]])
            else:
                v_shaped = self.v_template + del_v + torch.reshape(torch.matmul(beta, shapedirs_sel), [-1, self.size[0], self.size[1]])
        else:
            if del_v is None:
                v_shaped = self.v_template.unsqueeze(0)
            else:
                v_shaped = self.v_template + del_v 

        # 2. Infer shape-dependent joint locations.
        Jx = torch.matmul(v_shaped[:, :, 0], self.J_regressor)
        Jy = torch.matmul(v_shaped[:, :, 1], self.J_regressor)
        Jz = torch.matmul(v_shaped[:, :, 2], self.J_regressor)
        J = torch.stack([Jx, Jy, Jz], dim=2)

        # 3. Add pose blend shapes
        # N x 24 x 3 x 3
        if pose is None:
            Rs = torch.reshape( batch_rodrigues(torch.reshape(theta, [-1, 3])), [-1, 35, 3, 3])
        else:
            Rs = pose
        # Ignore global rotation.
        pose_feature = torch.reshape(Rs[:, 1:, :, :] - torch.eye(3).to(device=device), [-1, 306])

        v_posed = torch.reshape(
            torch.matmul(pose_feature, self.posedirs),
            [-1, self.size[0], self.size[1]]) + v_shaped

        #-------------------------
        # new: add corrections of bone lengths to the template  (before hypothetical pose blend shapes!)
        # see biggs batch_lbs.py
        betas_scale = torch.exp(betas_logscale @ self.betas_scale_mask.to(betas_logscale.device))
        scaling_factors = betas_scale.reshape(-1, 35, 3)
        scale_factors_3x3 = torch.diag_embed(scaling_factors, dim1=-2, dim2=-1)

        # 4. Get the global joint location
        # self.J_transformed, A = batch_global_rigid_transformation(Rs, J, self.parents)
        self.J_transformed, A = batch_global_rigid_transformation_biggs(Rs, J, self.parents, scale_factors_3x3, betas_logscale=betas_logscale)

        # 2-BONES. Calculate bone lengths
        all_bone_lengths = self._caclulate_bone_lengths_from_J(J, betas_logscale)
        # selected_bone_lengths = all_bone_lengths[:, IDXS_BONES_NO_REDUNDANCY]
        #-------------------------

        # 5. Do skinning:
        num_batch = Rs.shape[0]
        
        weights_t = self.weights.repeat([num_batch, 1])
        W = torch.reshape(weights_t, [num_batch, -1, 35])

            
        T = torch.reshape(
            torch.matmul(W, torch.reshape(A, [num_batch, 35, 16])),
                [num_batch, -1, 4, 4])
        v_posed_homo = torch.cat(
                [v_posed, torch.ones([num_batch, v_posed.shape[1], 1]).to(device=device)], 2)
        v_homo = torch.matmul(T, v_posed_homo.unsqueeze(-1))

        verts = v_homo[:, :, :3, 0]

        if trans is None:
            trans = torch.zeros((num_batch,3)).to(device=device)

        verts = verts + trans[:,None,:]

        # Get joints:
        joint_x = torch.matmul(verts[:, :, 0], self.J_regressor)
        joint_y = torch.matmul(verts[:, :, 1], self.J_regressor)
        joint_z = torch.matmul(verts[:, :, 2], self.J_regressor)
        joints = torch.stack([joint_x, joint_y, joint_z], dim=2)

        # New... (see https://github.com/benjiebob/SMALify/blob/master/smal_model/smal_torch.py)
        joints = torch.cat([
            joints,
            verts[:, None, 1863], # end_of_nose
            verts[:, None, 26], # chin
            verts[:, None, 2124], # right ear tip
            verts[:, None, 150], # left ear tip
            verts[:, None, 3055], # left eye
            verts[:, None, 1097], # right eye
            ], dim = 1) 

        if keyp_conf == 'blue' or keyp_conf == 'dict':
            # Generate keypoints
            nLandmarks = KEY_VIDS.shape[0]      # 24
            j3d = torch.zeros((num_batch, nLandmarks, 3)).to(device=device)
            for j in range(nLandmarks):
                j3d[:, j,:] = torch.mean(verts[:, KEY_VIDS[j],:], dim=1)  # translation is already added to the vertices
            joints_blue = j3d

        joints_red = joints[:, :-6, :]
        joints_green = joints[:, CANONICAL_MODEL_JOINTS, :]

        if keyp_conf == 'red':
            relevant_joints = joints_red
        elif keyp_conf == 'green':
            relevant_joints = joints_green
        elif keyp_conf == 'blue':
            relevant_joints = joints_blue
        elif keyp_conf == 'dict':
            relevant_joints = {'red': joints_red,
                            'green': joints_green,
                            'blue': joints_blue}
        else:
            raise NotImplementedError

        if get_all_info:
            return verts, relevant_joints, Rs, all_bone_lengths
        else:
            if get_skin:
                return verts, relevant_joints, Rs        # , v_shaped
            else:
                return relevant_joints