Kinematic dataset VAE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "from torch import nn \n",
    "import torch; torch.manual_seed(42)\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils\n",
    "import torch.distributions\n",
    "import torchvision\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt; plt.rcParams['figure.dpi'] = 200\n",
    "import pandas as pd\n",
    "from os import listdir, path\n",
    "import glob\n",
    "from sklearn.model_selection import train_test_split\n",
    "import torch.optim as optim\n",
    "import torch.distributions as dist\n",
    "import os "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = 'bvh_to_csv'\n",
    "data_files = glob.glob(data_dir + '/*.csv')\n",
    "tmps = listdir(data_dir)\n",
    "ems = []\n",
    "for i in range(len(tmps)):\n",
    "    tmp = tmps[i].split('.')\n",
    "    if tmp[0][4].isalpha():\n",
    "        emotion = tmp[0][3:5]\n",
    "    else:\n",
    "        emotion = tmp[0][3]\n",
    "    ems.append(emotion)\n",
    "    \n",
    "mapping = {\n",
    "    'A': 0,\n",
    "    'D': 1,\n",
    "    'F': 2,\n",
    "    'H': 3,\n",
    "    'N': 4,\n",
    "    'SA': 5,\n",
    "    'SU': 6\n",
    "}\n",
    "\n",
    "# Use list comprehension to convert elements to numbers\n",
    "labels = [mapping[item] for item in ems]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1401"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1989\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "# Define the path to the directory containing the mocap data files\n",
    "data_dir = 'bvh_to_csv'\n",
    "\n",
    "# List all the mocap data files\n",
    "data_files = glob.glob(data_dir + '/*.csv')\n",
    "\n",
    "# Initialize empty lists to store the data\n",
    "mocap_data = []\n",
    "emotion_labels = []\n",
    "length = []\n",
    "# Load the data from each file\n",
    "for file in data_files:\n",
    "    # Load the file and extract the coordinates\n",
    "    file_data = np.genfromtxt(file, delimiter=',')\n",
    "\n",
    "\n",
    "    # print(file_data.shape)\n",
    "   \n",
    "    # print(max_length)\n",
    "\n",
    "    # padded_mocap_data = []\n",
    "    # for sequence in mocap_data:\n",
    "    #     num_frames_to_pad = max_length - sequence.shape[0]\n",
    "    #     last_row = sequence[-1]  # Get the last row\n",
    "    #     padded_sequence = np.concatenate((sequence, np.tile(last_row, (num_frames_to_pad, 1))), axis=0)\n",
    "    #     padded_mocap_data.append(padded_sequence)\n",
    "\n",
    "    # file_data = pd.read_csv(file, sep=',', skiprows=1)\n",
    "    # file_data = file_data.iloc[:,1:]\n",
    "# print(data)\n",
    "    \n",
    "    min_val = np.min(file_data)\n",
    "    normalized_data = file_data - min_val\n",
    "    range_val = np.max(file_data) - min_val\n",
    "    normalized_data /= range_val\n",
    "    file_data = normalized_data\n",
    "    \n",
    "#     # Append the coordinates to the mocap_data list\n",
    "    mocap_data.append(file_data)\n",
    "\n",
    "\n",
    "max_length = max([sequence.shape[0] for sequence in mocap_data])\n",
    "\n",
    "padded_mocap_data = []\n",
    "for sequence in mocap_data:\n",
    "    num_frames_to_pad = max_length - sequence.shape[0]\n",
    "    last_row = sequence[-1]  # Get the last row\n",
    "    padded_sequence = np.concatenate((sequence, np.tile(last_row, (num_frames_to_pad, 1))), axis=0)\n",
    "    padded_mocap_data.append(padded_sequence)\n",
    "# Convert the emotion labels to a NumPy array\n",
    "emotion_label = np.array(labels)\n",
    "\n",
    "# Perform any necessary preprocessing steps on the mocap data and emotion labels\n",
    "# For example, you can normalize the data or encode the emotion labels\n",
    "\n",
    "# Split the data into training and testing sets\n",
    "X_train, X_test, y_train, y_test = train_test_split(padded_mocap_data, emotion_label, test_size=0.2, random_state=42)\n",
    "print(max_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kl_divergence_loss(mu, logvar):\n",
    "    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n",
    "    return kl_loss\n",
    "\n",
    "def reconstruction_loss(input,target):\n",
    "    recon_loss = nn.CrossEntropyLoss()(input,target)\n",
    "    return recon_loss\n",
    "\n",
    "def total_loss( mu, logvar,input,target):\n",
    "    kl_loss = kl_divergence_loss(mu, logvar)\n",
    "    recon_loss = reconstruction_loss(input,target)\n",
    "\n",
    "    total_loss = 0.005 * kl_loss + 0.95* recon_loss\n",
    "    return total_loss\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BatchFlatten(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BatchFlatten,self).__init__()\n",
    "        self._name = 'batch_flatten'\n",
    "    def forward(self, x):\n",
    "        return x.view(x.shape[0],-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NormalDistDecoder(nn.Module):\n",
    "    def __init__(self, num_feat_in, latentD):\n",
    "        super(NormalDistDecoder, self).__init__()\n",
    "\n",
    "        self.mu = nn.Linear(num_feat_in, latentD)\n",
    "        self.logvar = nn.Linear(num_feat_in, latentD)\n",
    "\n",
    "    def forward(self, Xout):\n",
    "        mu_logvar =torch.distributions.normal.Normal(self.mu(Xout), F.softplus(self.logvar(Xout)))\n",
    "        return mu_logvar\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VPoser(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(VPoser, self).__init__()\n",
    "\n",
    "        num_neurons, self.latentD = 512,512\n",
    "        nb_frames = 1989\n",
    "\n",
    "        self.num_joints = 58\n",
    "        n_features = self.num_joints * 3 * nb_frames\n",
    "\n",
    "        self.encoder_net = nn.Sequential(\n",
    "            BatchFlatten(),\n",
    "            nn.BatchNorm1d(n_features),\n",
    "            nn.Linear(n_features, num_neurons),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.BatchNorm1d(num_neurons),\n",
    "            nn.Dropout(0.1),\n",
    "            nn.Linear(num_neurons, num_neurons),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(num_neurons, self.latentD),\n",
    "            # NormalDistDecoder(num_neurons, self.latentD)\n",
    "            \n",
    "            \n",
    "            \n",
    "        )\n",
    "\n",
    "        self.decoder_net = nn.Sequential(\n",
    "            nn.Linear(self.latentD, num_neurons),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Dropout(0.1),\n",
    "            nn.Linear(num_neurons, num_neurons),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(num_neurons, n_features),\n",
    "            # ContinousRotReprDecoder(),\n",
    "        )\n",
    "        self.mu_net = nn.Linear(512, 512)\n",
    "        self.logvar_net = nn.Linear(512,512)\n",
    "        # self.decoder_net = NormalDistDecoder(self.latentD, 512)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    def encode(self, pose_body):\n",
    "        '''\n",
    "        :param Pin: Nx(numjoints*3)\n",
    "        :param rep_type: 'matrot'/'aa' for matrix rotations or axis-angle\n",
    "        :return:\n",
    "        \n",
    "        '''\n",
    "        # print('posebody', pose_body.shape)\n",
    "        # self.mu_net = nn.Linear(32, 512)\n",
    "        # self.logvar_net = nn.Linear(32,512)\n",
    "        q_z = self.encoder_net(pose_body)\n",
    "        # print(\"qz\",q_z.shape)\n",
    "        mu = self.mu_net(q_z)\n",
    "        logvar = self.logvar_net(q_z)\n",
    "        return q_z, mu, logvar\n",
    "\n",
    "    def decode(self, Zin):\n",
    "        # print('zin' ,Zin.shape)\n",
    "        bs = Zin.shape[0]\n",
    "        # print(bs)\n",
    "\n",
    "        prec = self.decoder_net(Zin)\n",
    "        prec1 = prec.view(bs,-1,174)\n",
    "        # print('prec' ,prec1.shape)\n",
    "\n",
    "        return {\n",
    "            # 'pose_body': matrot2aa(prec.view(-1, 3, 3)).view(bs, -1, 3),\n",
    "            'pose_body_matrot': prec.view(bs,-1,174)\n",
    "        }\n",
    "\n",
    "\n",
    "    def forward(self, pose_body):    \n",
    "\n",
    "        q_z, mu, logvar = self.encode(pose_body)\n",
    "        # print(q_z.shape)\n",
    "        # q_z_sample = q_z.rsample()\n",
    "        q_z_sample = self.reparameterize(mu, logvar)\n",
    "        # print('qz sample ', q_z_sample.shape)\n",
    "\n",
    "        decode_results = self.decode(q_z_sample)\n",
    "        pose_body_decoded = decode_results['pose_body_matrot']  # Extract the decoded pose tensor\n",
    "\n",
    "        # decode_results.update({'poZ_body_mean': q_z.mean, 'poZ_body_std': q_z.scale, 'q_z': q_z})\n",
    "        return pose_body_decoded, mu , logvar\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps * std\n",
    "    \n",
    "    \n",
    "    \n",
    "    def sample_poses(self, num_poses, seed=None):\n",
    "        np.random.seed(seed)\n",
    "\n",
    "        some_weight = [a for a in self.parameters()][0]\n",
    "        dtype = some_weight.dtype\n",
    "        device = some_weight.device\n",
    "        self.eval()\n",
    "        with torch.no_grad():\n",
    "            Zgen = torch.tensor(np.random.normal(0., 1., size=(num_poses, self.latentD)), dtype=dtype, device=device)\n",
    "\n",
    "        return self.decode(Zgen)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "# Convert the mocap data and emotion labels to NumPy arrays\n",
    "X_train_array = np.array(X_train, dtype=np.float32)\n",
    "y_train_array = np.array(y_train, dtype=np.compat.long)\n",
    "X_test_array = np.array(X_test, dtype=np.float32)\n",
    "y_test_array = np.array(y_test, dtype=np.compat.long)\n",
    "\n",
    "# Convert NumPy arrays to PyTorch tensors\n",
    "X_train_tensor = torch.tensor(X_train_array)\n",
    "y_train_tensor = torch.tensor(y_train_array)\n",
    "X_test_tensor = torch.tensor(X_test_array)\n",
    "y_test_tensor = torch.tensor(y_test_array)\n",
    "\n",
    "# Create a custom dataset class to load the data\n",
    "class MocapDataset(Dataset):\n",
    "    def __init__(self, data, labels):\n",
    "        self.data = data\n",
    "        self.labels = labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx], self.labels[idx]\n",
    "\n",
    "# Create DataLoader for training and testing sets\n",
    "train_dataset = MocapDataset(X_train_tensor, y_train_tensor)\n",
    "train_loader = DataLoader(train_dataset, batch_size=128, shuffle=False)\n",
    "\n",
    "test_dataset = MocapDataset(X_test_tensor, y_test_tensor)\n",
    "test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean Squared Error: 0.2621057331562042\n",
      "Epoch 1/10, Loss: 10.963930130004883\n",
      "Mean Squared Error: 0.2623075842857361\n",
      "Epoch 2/10, Loss: 10.983806610107422\n",
      "Mean Squared Error: 0.2622583210468292\n",
      "Epoch 3/10, Loss: 10.903276443481445\n",
      "Mean Squared Error: 0.2621668577194214\n",
      "Epoch 4/10, Loss: 11.008668899536133\n",
      "Mean Squared Error: 0.2619835138320923\n",
      "Epoch 5/10, Loss: 10.939143180847168\n",
      "Mean Squared Error: 0.2619412839412689\n",
      "Epoch 6/10, Loss: 10.955436706542969\n",
      "Mean Squared Error: 0.2619992792606354\n",
      "Epoch 7/10, Loss: 10.964600563049316\n",
      "Mean Squared Error: 0.2619350552558899\n",
      "Epoch 8/10, Loss: 10.923131942749023\n",
      "Mean Squared Error: 0.2621893882751465\n",
      "Epoch 9/10, Loss: 10.911256790161133\n",
      "Mean Squared Error: 0.2622235119342804\n",
      "Epoch 10/10, Loss: 11.017494201660156\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "# Define the VPoser model\n",
    "model = VPoser()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.000001)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "# Training loop\n",
    "num_epochs = 10\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    for batch_data, batch_labels in train_loader:\n",
    "        # print(batch_data[0][0])\n",
    "        optimizer.zero_grad()\n",
    "        outputs, mu, logvar = model(batch_data)\n",
    "        # print(outputs[0][0])\n",
    "       \n",
    "        \n",
    "        labels = batch_labels.repeat_interleave(174)\n",
    "        # print(labels.shape)\n",
    "\n",
    "# Reshape the tensor to size (128, 3)\n",
    "        labels = labels.reshape(-1,174)\n",
    "# ENCODE THE LABELS AND DO A CLUSTERING  SEE THE DISTROBUTION OF LATENT \n",
    "        # print(labels.shape)\n",
    "        labels = labels.long()    \n",
    "        # loss = criterion(outputs, labels)\n",
    "        loss = total_loss(mu,logvar,outputs,labels) \n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "\n",
    "        mse_criterion = torch.nn.MSELoss()\n",
    "        mse = mse_criterion(outputs, batch_data)\n",
    "    print(\"Mean Squared Error:\", mse.item())\n",
    "\n",
    "      \n",
    "       \n",
    "\n",
    "        \n",
    "# LATENT PROBABILITY  \n",
    "    # Print the loss at the end of each epoch\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {loss}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Loss: 8.05859661102295 \n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "test_loss = 0.0\n",
    "num_samples = 0\n",
    "for batch_data, batch_labels in test_loader:\n",
    "    outputs, mu, logvar = model(batch_data)\n",
    "    labels = batch_labels.repeat_interleave(174)\n",
    "    labels = labels.reshape(-1,174)\n",
    "    labels = labels.long()    \n",
    "    # loss = criterion(outputs,labels)\n",
    "    loss = total_loss(mu,logvar,outputs,labels)\n",
    "    \n",
    "print(f\" Loss: {loss} \")\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "IMS",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}