models/motion_representation.py

import random
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import smplx
import copy
from .motion_encoder import * 

# ----------- AE, VAE ------------- #
class VAEConvZero(nn.Module):
    def __init__(self, args):
        super(VAEConvZero, self).__init__()
        self.encoder = VQEncoderV5(args)
        # self.quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        self.decoder = VQDecoderV5(args)
        
    def forward(self, inputs):
        pre_latent = self.encoder(inputs)
        # print(pre_latent.shape)
        # embedding_loss, vq_latent, _, perplexity = self.quantizer(pre_latent)
        rec_pose = self.decoder(pre_latent)
        return {
            # "poses_feat":vq_latent,
            # "embedding_loss":embedding_loss,
            # "perplexity":perplexity,
            "rec_pose": rec_pose
            }
    
class VAEConv(nn.Module):
    def __init__(self, args):
        super(VAEConv, self).__init__()
        self.encoder = VQEncoderV3(args)
        self.decoder = VQDecoderV3(args)
        self.fc_mu = nn.Linear(args.vae_length, args.vae_length)
        self.fc_logvar = nn.Linear(args.vae_length, args.vae_length)
        self.variational = args.variational
        
    def forward(self, inputs):
        pre_latent = self.encoder(inputs)
        mu, logvar = None, None
        if self.variational:
            mu = self.fc_mu(pre_latent)
            logvar = self.fc_logvar(pre_latent)
            pre_latent = reparameterize(mu, logvar)
        rec_pose = self.decoder(pre_latent)
        return {
            "poses_feat":pre_latent,
            "rec_pose": rec_pose,
            "pose_mu": mu,
            "pose_logvar": logvar,
            }
    
    def map2latent(self, inputs):
        pre_latent = self.encoder(inputs)
        if self.variational:
            mu = self.fc_mu(pre_latent)
            logvar = self.fc_logvar(pre_latent)
            pre_latent = reparameterize(mu, logvar)
        return pre_latent
    
    def decode(self, pre_latent):
        rec_pose = self.decoder(pre_latent)
        return rec_pose

class VAESKConv(VAEConv):
    def __init__(self, args):
        super(VAESKConv, self).__init__(args)
        smpl_fname = args.data_path_1+'smplx_models/smplx/SMPLX_NEUTRAL_2020.npz'
        smpl_data = np.load(smpl_fname, encoding='latin1')
        parents = smpl_data['kintree_table'][0].astype(np.int32)
        edges = build_edge_topology(parents)
        self.encoder = LocalEncoder(args, edges)
        self.decoder = VQDecoderV3(args)
        
class VAEConvMLP(VAEConv):
    def __init__(self, args):
        super(VAEConvMLP, self).__init__(args)
        self.encoder = PoseEncoderConv(args.vae_test_len, args.vae_test_dim, feature_length=args.vae_length)
        self.decoder = PoseDecoderConv(args.vae_test_len, args.vae_test_dim, feature_length=args.vae_length)
 
class VAELSTM(VAEConv):
    def __init__(self, args):
        super(VAELSTM, self).__init__(args)
        pose_dim = args.vae_test_dim
        feature_length = args.vae_length
        self.encoder = PoseEncoderLSTM_Resnet(pose_dim, feature_length=feature_length)
        self.decoder = PoseDecoderLSTM(pose_dim, feature_length=feature_length)

class VAETransformer(VAEConv):
    def __init__(self, args):
        super(VAETransformer, self).__init__(args)
        self.encoder = Encoder_TRANSFORMER(args)
        self.decoder = Decoder_TRANSFORMER(args)

# ----------- VQVAE --------------- #
class VQVAEConv(nn.Module):
    def __init__(self, args):
        super(VQVAEConv, self).__init__()
        self.encoder = VQEncoderV3(args)
        self.quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        self.decoder = VQDecoderV3(args)
        
    def forward(self, inputs):
        pre_latent = self.encoder(inputs)
        # print(pre_latent.shape)
        embedding_loss, vq_latent, _, perplexity = self.quantizer(pre_latent)
        rec_pose = self.decoder(vq_latent)
        return {
            "poses_feat":vq_latent,
            "embedding_loss":embedding_loss,
            "perplexity":perplexity,
            "rec_pose": rec_pose
            }
    
    def map2index(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        return index
    
    def map2latent(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        z_q = self.quantizer.get_codebook_entry(index)
        return z_q
    
    def decode(self, index):
        z_q = self.quantizer.get_codebook_entry(index)
        rec_pose = self.decoder(z_q)
        return rec_pose

class VQVAESKConv(VQVAEConv):
    def __init__(self, args):
        super(VQVAESKConv, self).__init__(args)
        smpl_fname = args.data_path_1+'smplx_models/smplx/SMPLX_NEUTRAL_2020.npz'
        smpl_data = np.load(smpl_fname, encoding='latin1')
        parents = smpl_data['kintree_table'][0].astype(np.int32)
        edges = build_edge_topology(parents)
        self.encoder = LocalEncoder(args, edges)


class VQVAEConvStride(nn.Module):
    def __init__(self, args):
        super(VQVAEConvStride, self).__init__()
        self.encoder = VQEncoderV4(args)
        self.quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        self.decoder = VQDecoderV4(args)
        
    def forward(self, inputs):
        pre_latent = self.encoder(inputs)
        # print(pre_latent.shape)
        embedding_loss, vq_latent, _, perplexity = self.quantizer(pre_latent)
        rec_pose = self.decoder(vq_latent)
        return {
            "poses_feat":vq_latent,
            "embedding_loss":embedding_loss,
            "perplexity":perplexity,
            "rec_pose": rec_pose
            }
    
    def map2index(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        return index
    
    def map2latent(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        z_q = self.quantizer.get_codebook_entry(index)
        return z_q
    
    def decode(self, index):
        z_q = self.quantizer.get_codebook_entry(index)
        rec_pose = self.decoder(z_q)
        return rec_pose

class VQVAEConvZero(nn.Module):
    def __init__(self, args):
        super(VQVAEConvZero, self).__init__()
        self.encoder = VQEncoderV5(args)
        self.quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        self.decoder = VQDecoderV5(args)
        
    def forward(self, inputs):
        pre_latent = self.encoder(inputs)
        # print(pre_latent.shape)
        embedding_loss, vq_latent, _, perplexity = self.quantizer(pre_latent)
        rec_pose = self.decoder(vq_latent)
        return {
            "poses_feat":vq_latent,
            "embedding_loss":embedding_loss,
            "perplexity":perplexity,
            "rec_pose": rec_pose
            }
    
    def map2index(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        return index
    
    def map2latent(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        z_q = self.quantizer.get_codebook_entry(index)
        return z_q
    
    def decode(self, index):
        z_q = self.quantizer.get_codebook_entry(index)
        rec_pose = self.decoder(z_q)
        return rec_pose
    

class VAEConvZero(nn.Module):
    def __init__(self, args):
        super(VAEConvZero, self).__init__()
        self.encoder = VQEncoderV5(args)
        # self.quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        self.decoder = VQDecoderV5(args)
        
    def forward(self, inputs):
        pre_latent = self.encoder(inputs)
        # print(pre_latent.shape)
        # embedding_loss, vq_latent, _, perplexity = self.quantizer(pre_latent)
        rec_pose = self.decoder(pre_latent)
        return {
            # "poses_feat":vq_latent,
            # "embedding_loss":embedding_loss,
            # "perplexity":perplexity,
            "rec_pose": rec_pose
            }
    
    # def map2index(self, inputs):
    #     pre_latent = self.encoder(inputs)
    #     index = self.quantizer.map2index(pre_latent)
    #     return index
    
    # def map2latent(self, inputs):
    #     pre_latent = self.encoder(inputs)
    #     index = self.quantizer.map2index(pre_latent)
    #     z_q = self.quantizer.get_codebook_entry(index)
    #     return z_q
    
    # def decode(self, index):
    #     z_q = self.quantizer.get_codebook_entry(index)
    #     rec_pose = self.decoder(z_q)
    #     return rec_pose


class VQVAEConvZero3(nn.Module):
    def __init__(self, args):
        super(VQVAEConvZero3, self).__init__()
        self.encoder = VQEncoderV5(args)
        self.quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        self.decoder = VQDecoderV5(args)
        
    def forward(self, inputs):
        pre_latent = self.encoder(inputs)
        # print(pre_latent.shape)
        embedding_loss, vq_latent, _, perplexity = self.quantizer(pre_latent)
        rec_pose = self.decoder(vq_latent)
        return {
            "poses_feat":vq_latent,
            "embedding_loss":embedding_loss,
            "perplexity":perplexity,
            "rec_pose": rec_pose
            }
    
    def map2index(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        return index
    
    def map2latent(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        z_q = self.quantizer.get_codebook_entry(index)
        return z_q
    
    def decode(self, index):
        z_q = self.quantizer.get_codebook_entry(index)
        rec_pose = self.decoder(z_q)
        return rec_pose

class VQVAEConvZero2(nn.Module):
    def __init__(self, args):
        super(VQVAEConvZero2, self).__init__()
        self.encoder = VQEncoderV5(args)
        self.quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        self.decoder = VQDecoderV7(args)
        
    def forward(self, inputs):
        pre_latent = self.encoder(inputs)
        # print(pre_latent.shape)
        embedding_loss, vq_latent, _, perplexity = self.quantizer(pre_latent)
        rec_pose = self.decoder(vq_latent)
        return {
            "poses_feat":vq_latent,
            "embedding_loss":embedding_loss,
            "perplexity":perplexity,
            "rec_pose": rec_pose
            }
    
    def map2index(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        return index
    
    def map2latent(self, inputs):
        pre_latent = self.encoder(inputs)
        index = self.quantizer.map2index(pre_latent)
        z_q = self.quantizer.get_codebook_entry(index)
        return z_q
    
    def decode(self, index):
        z_q = self.quantizer.get_codebook_entry(index)
        rec_pose = self.decoder(z_q)
        return rec_pose

class VQVAE2(nn.Module):
    def __init__(self, args):
        super(VQVAE2, self).__init__()
        # Bottom-level encoder and decoder
        args_bottom = copy.deepcopy(args)
        args_bottom.vae_layer = 2
        self.bottom_encoder = VQEncoderV6(args_bottom)
        self.bottom_quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        args_bottom.vae_test_dim = args.vae_test_dim
        self.bottom_decoder = VQDecoderV6(args_bottom)
        
        # Top-level encoder and decoder
        args_top = copy.deepcopy(args)
        args_top.vae_layer = 3
        args_top.vae_test_dim = args.vae_length
        self.top_encoder = VQEncoderV3(args_top)  # Adjust according to the top level's design
        self.quantize_conv_t = nn.Conv1d(args.vae_length+args.vae_length, args.vae_length, 1)
        self.top_quantizer = Quantizer(args.vae_codebook_size, args.vae_length, args.vae_quantizer_lambda)
        # self.upsample_t_up = nn.Upsample(scale_factor=2, mode='nearest')
        layers = [
                nn.Upsample(scale_factor=2, mode='nearest'),
                nn.Conv1d(args.vae_length, args.vae_length, kernel_size=3, stride=1, padding=1),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Upsample(scale_factor=2, mode='nearest'),
                nn.Conv1d(args.vae_length, args.vae_length, kernel_size=3, stride=1, padding=1),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Upsample(scale_factor=2, mode='nearest'),
                nn.Conv1d(args.vae_length, args.vae_length, kernel_size=3, stride=1, padding=1),
                nn.LeakyReLU(0.2, inplace=True)
            ]
        self.upsample_t= nn.Sequential(*layers)
        self.top_decoder = VQDecoderV3(args_top)  # Adjust to handle top level features appropriately

    def forward(self, inputs):
        # Bottom-level processing
        enc_b = self.bottom_encoder(inputs)
        enc_t = self.top_encoder(enc_b)
        #print(enc_b.shape, enc_t.shape)
        top_embedding_loss, quant_t, _, top_perplexity = self.top_quantizer(enc_t)
        #print(quant_t.shape)
        dec_t = self.top_decoder(quant_t)
        #print(dec_t.shape)
        enc_b = torch.cat([dec_t, enc_b], dim=2).permute(0,2,1)
        #print(enc_b.shape)
        quant_b = self.quantize_conv_t(enc_b).permute(0,2,1)
        #print("5",quant_b.shape)
        bottom_embedding_loss, quant_b, _, bottom_perplexity = self.bottom_quantizer(quant_b)
        #print("6",quant_b.shape)
        upsample_t = self.upsample_t(quant_t.permute(0,2,1)).permute(0,2,1)
        #print("7",upsample_t.shape)
        quant = torch.cat([upsample_t, quant_b], 2)
        rec_pose = self.bottom_decoder(quant)
        # print(quant_t.shape, quant_b.shape, rec_pose.shape)
        return {
            "poses_feat_top": quant_t,
            "pose_feat_bottom": quant_b,
            "embedding_loss":top_embedding_loss+bottom_embedding_loss,
            #"perplexity":perplexity,
            "rec_pose": rec_pose
            }
    
    def map2index(self, inputs):
        enc_b = self.bottom_encoder(inputs)
        enc_t = self.top_encoder(enc_b)
        
        _, quant_t, _, _ = self.top_quantizer(enc_t)
        top_index = self.top_quantizer.map2index(enc_t)
        dec_t = self.top_decoder(quant_t)

        enc_b = torch.cat([dec_t, enc_b], dim=2).permute(0,2,1)
        #print(enc_b.shape)
        quant_b = self.quantize_conv_t(enc_b).permute(0,2,1)
        # quant_b = self.quantize_conv_t(enc_b)
        bottom_index = self.bottom_quantizer.map2index(quant_b)
        return top_index, bottom_index
    
    def get_top_laent(self, top_index):
        z_q_top = self.top_quantizer.get_codebook_entry(top_index)
        return z_q_top
    
    def map2latent(self, inputs):
        enc_b = self.bottom_encoder(inputs)
        enc_t = self.top_encoder(enc_b)
        
        _, quant_t, _, _ = self.top_quantizer(enc_t)
        top_index = self.top_quantizer.map2index(enc_t)
        dec_t = self.top_decoder(quant_t)

        enc_b = torch.cat([dec_t, enc_b], dim=2).permute(0,2,1)
        #print(enc_b.shape)
        quant_b = self.quantize_conv_t(enc_b).permute(0,2,1)
        # quant_b = self.quantize_conv_t(enc_b)
        bottom_index = self.bottom_quantizer.map2index(quant_b)
        z_q_top = self.top_quantizer.get_codebook_entry(top_index)
        z_q_bottom = self.bottom_quantizer.get_codebook_entry(bottom_index)
        return z_q_top, z_q_bottom
    
    def map2latent_top(self, inputs):
        enc_b = self.bottom_encoder(inputs)
        enc_t = self.top_encoder(enc_b)
        top_index = self.top_quantizer.map2index(enc_t)
        z_q_top = self.top_quantizer.get_codebook_entry(top_index)
        return z_q_top
    
    def decode(self, top_index, bottom_index):
        quant_t = self.top_quantizer.get_codebook_entry(top_index)
        quant_b = self.bottom_quantizer.get_codebook_entry(bottom_index)
        upsample_t = self.upsample_t(quant_t.permute(0,2,1)).permute(0,2,1)
        #print("7",upsample_t.shape)
        quant = torch.cat([upsample_t, quant_b], 2)
        rec_pose = self.bottom_decoder(quant)      
        return rec_pose