BERT.KeyFramesFilling.Ubidata.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\"\n",
    "\n",
    "import joblib\n",
    "import plotly.graph_objects as go\n",
    "import numpy as np\n",
    "import random\n",
    "import torch\n",
    "import transformers\n",
    "# Set random seed for reproducibility\n",
    "manualSeed = 999\n",
    "#manualSeed = random.randint(1, 10000) # use if you want new results\n",
    "print(\"Random Seed: \", manualSeed)\n",
    "random.seed(manualSeed)\n",
    "torch.manual_seed(manualSeed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use to save results in .bvh file\n",
    "\n",
    "import sys\n",
    "sys.path.append('utils/PyMO/pymo/')\n",
    "\n",
    "import sys\n",
    "sys.path.append('utils/Poses/PyMO/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.load('data/ubisoft.rotation_matrix.300.npy', allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_train = np.array([i for i in data if i['type'] == 'train'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_normalization = {}\n",
    "\n",
    "for i in range(198, 201):\n",
    "    tmp = np.concatenate([j['pose'] for j in data], 0)[:, i]\n",
    "    tmp_mean = tmp.mean()\n",
    "    tmp_std = tmp.std()\n",
    "    tmp_min = tmp.min()\n",
    "    tmp_max = tmp.max()\n",
    "\n",
    "    if data_normalization.get('pose', None) is None:\n",
    "        data_normalization['pose'] = {}\n",
    "    if data_normalization['pose'].get(i, None) is None:\n",
    "        data_normalization['pose'][i] = {}\n",
    "    data_normalization['pose'][i]['mean'] = tmp_mean\n",
    "    data_normalization['pose'][i]['std'] = tmp_std\n",
    "    data_normalization['pose'][i]['min'] = tmp_min\n",
    "    data_normalization['pose'][i]['max'] = tmp_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def norm_train_poses(sample):\n",
    "    for i in range(198, 201):\n",
    "        sample[:, :, i] = (sample[:, :, i] - data_normalization['pose'][i]['min']) / (data_normalization['pose'][i]['max'] - data_normalization['pose'][i]['min'])\n",
    "\n",
    "    return sample\n",
    "\n",
    "def norm_train_music(sample):\n",
    "    for i in range(35):\n",
    "        sample[:, :, i] = (sample[:, :, i] - data_normalization['music'][i]['min']) / (data_normalization['music'][i]['max'] - data_normalization['music'][i]['min'])\n",
    "    \n",
    "    return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def unnorm_train_poses(sample):\n",
    "    for i in range(198, 201):\n",
    "            sample[:, :, i] = sample[:, :, i] * (data_normalization['pose'][i]['max'] - data_normalization['pose'][i]['min']) + data_normalization['pose'][i]['min']\n",
    "            \n",
    "    return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.ipc_collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda:0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_config = transformers.ConvBertConfig().to_dict()\n",
    "config = transformers.PretrainedConfig().to_dict()\n",
    "\n",
    "bert_config['num_hidden_layers'] = 10\n",
    "bert_config['num_attention_heads'] = 8\n",
    "bert_config['output_hidden_states'] = True\n",
    "bert_config['hidden_size'] = 512\n",
    "bert_config['embedding_size'] = 512\n",
    "bert_config['max_position_embeddings'] = 300\n",
    "# bert_config['type_vocab_size'] = 24\n",
    "\n",
    "bert_config = transformers.ConvBertConfig().from_dict(bert_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PoseBERT(torch.nn.Module):\n",
    "    def __init__(self, bert_config):\n",
    "        super(PoseBERT, self).__init__()\n",
    "        self.bert = transformers.ConvBertModel(bert_config)\n",
    "\n",
    "        self.conv1 = torch.nn.Conv1d(512, 256, 3, padding=(1,))\n",
    "        self.conv2 = torch.nn.Conv1d(256, 201, 1)\n",
    "        self.norm = torch.nn.LayerNorm((300, 201))\n",
    "#         self.linear = torch.nn.Linear(512, 201)\n",
    "        self.activation = torch.nn.Tanh()\n",
    "        \n",
    "    def single_pass(self, xp, attention_mask):\n",
    "        xp = torch.nn.functional.pad(xp, (0, 512 - xp.shape[-1]))\n",
    "        \n",
    "        poses_output_seed = self.bert(inputs_embeds=xp, attention_mask=attention_mask)['last_hidden_state']\n",
    "        poses_output_seed = poses_output_seed.transpose(1, -1)\n",
    "        decoder_output_pose = self.conv2(self.conv1(poses_output_seed))\n",
    "        decoder_output_pose = self.norm(decoder_output_pose.transpose(1, -1))\n",
    "\n",
    "        return None, decoder_output_pose, None\n",
    "    \n",
    "    def forward(self, input_xp, attention_mask, device='cpu', noise=None, train=True):\n",
    "        if train:\n",
    "            _, predict, _ = self.single_pass(input_xp, attention_mask)\n",
    "                \n",
    "            return predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "poseBert = PoseBERT(bert_config).to(device)\n",
    "# discriminate = DiscriminatorSingle(context_window=5).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# poseBert.load_state_dict(torch.load('/ess_storage/storage/home/skhn2/Poses/models/G.40/generator.step36427.G0.066.pt'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# history arrays\n",
    "G_losses = []\n",
    "D_losses = []\n",
    "\n",
    "iters = 0\n",
    "epochs = 1000\n",
    "batch_size = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# criterion = FrechetDistance()\n",
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "lrG = 1e-4\n",
    "beta1 = 0.9\n",
    "\n",
    "# fillers\n",
    "real_label = 1.\n",
    "fake_label = 0.\n",
    "\n",
    "optimizerG = torch.optim.Adam(poseBert.parameters(), lr=lrG, betas=(beta1, 0.999))\n",
    "# optimizerD = torch.optim.AdamW(discriminate.parameters(), lr=lrD, betas=(beta1, 0.999))\n",
    "\n",
    "# scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizerG, max_lr=1e-5, steps_per_epoch=len(data_cls)//batch_size, epochs=500, final_div_factor=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_train = data_train[:len(data_train) // batch_size * batch_size]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames_seed = 300\n",
    "frames_music = 300\n",
    "frames_min_length = 120\n",
    "\n",
    "masked_probability = 0.25\n",
    "zero_probability = 0.80\n",
    "replace_probability = 0.10\n",
    "same_probability = 0.10\n",
    "\n",
    "# frames_seed = 300\n",
    "# frames_music = 300\n",
    "# frames_min_length = 120\n",
    "\n",
    "# masked_probability = 0.001\n",
    "# zero_probability = 1.0\n",
    "# replace_probability = 0.0\n",
    "# same_probability = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def generate_masked_batch(pose_batch):\n",
    "    masked_pose_batch = torch.tensor(pose_batch)\n",
    "    original_pose_batch = torch.tensor(pose_batch)\n",
    "\n",
    "    samples_lengths = np.random.randint(frames_min_length, frames_seed, (pose_batch.shape[0]))\n",
    "    masks = np.ones((pose_batch.shape[0], frames_seed))\n",
    "    attention_mask = np.ones((pose_batch.shape[0], frames_seed))\n",
    "\n",
    "    for j, sample_length in enumerate(samples_lengths):\n",
    "        attention_mask[j][sample_length:] *= 0\n",
    "        mask = np.random.choice([1, 0], p=[1-masked_probability, masked_probability], size=(sample_length))\n",
    "        masks[j][:len(mask)] *= mask\n",
    "        mask_type = np.random.choice([0, 1, 2], p=[zero_probability, replace_probability, same_probability], size=(mask == 0).sum())\n",
    "        \n",
    "        masked_pose_batch[j][sample_length:] *= 0\n",
    "        for i, index in enumerate(np.where(mask == 0)[0]):\n",
    "            if mask_type[i] == 0:\n",
    "                masked_pose_batch[j][index] *= 0\n",
    "\n",
    "            elif mask_type[i] == 1:\n",
    "                select_sample_to_replace = np.random.randint(pose_batch.shape[0])\n",
    "                lst = np.arange(samples_lengths[select_sample_to_replace])\n",
    "                lst = lst[np.where(lst != index)]\n",
    "                random_replace = np.random.choice(lst)\n",
    "\n",
    "                masked_pose_batch[j][index] = original_pose_batch[select_sample_to_replace][random_replace]\n",
    "                \n",
    "    return masked_pose_batch, masks[:, :, None].repeat(pose_batch.shape[-1], -1), attention_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def generate_masked_batch_patched(pose_batch, pattern=0):\n",
    "    masked_pose_batch = torch.tensor(pose_batch)\n",
    "    original_pose_batch = torch.tensor(pose_batch)\n",
    "\n",
    "    samples_lengths = np.array([100 for _ in range(pose_batch.shape[0])])# np.random.randint(frames_min_length, frames_seed, (pose_batch.shape[0]))\n",
    "    masks = np.ones((pose_batch.shape[0], frames_seed))\n",
    "    attention_mask = np.ones((pose_batch.shape[0], frames_seed))\n",
    "\n",
    "    for j, sample_length in enumerate(samples_lengths):\n",
    "        attention_mask[j][sample_length:] *= 0\n",
    "        if pattern == 0:\n",
    "            mask = np.ones((sample_length))\n",
    "            mask[0] = 0\n",
    "        elif pattern == 1:\n",
    "            tmp = [1] * 5 + [0] * 10\n",
    "            mask = np.array((tmp * 30)[:sample_length])\n",
    "        elif pattern == 2:\n",
    "            mask = np.zeros((sample_length))\n",
    "            mask[::10] = 1\n",
    "        elif pattern == 3:\n",
    "            mask = np.ones((sample_length))\n",
    "            mask[88:98] = 0\n",
    "        \n",
    "#         mask = np.ones((sample_length))\n",
    "        \n",
    "#         mask[30:35] = 0\n",
    "#         mask[0] = 0\n",
    "#         mask = np.random.choice([1, 0], p=[1-masked_probability, masked_probability], size=(sample_length))\n",
    "        masks[j][:len(mask)] *= mask\n",
    "        mask_type = np.random.choice([0, 1, 2], p=[zero_probability, replace_probability, same_probability], size=(mask == 0).sum())\n",
    "        \n",
    "        masked_pose_batch[j][sample_length:] *= 0\n",
    "        for i, index in enumerate(np.where(mask == 0)[0]):\n",
    "            if mask_type[i] == 0:\n",
    "                masked_pose_batch[j][index] *= 0\n",
    "\n",
    "            elif mask_type[i] == 1:\n",
    "                select_sample_to_replace = np.random.randint(pose_batch.shape[0])\n",
    "                lst = np.arange(samples_lengths[select_sample_to_replace])\n",
    "                lst = lst[np.where(lst != index)]\n",
    "                random_replace = np.random.choice(lst)\n",
    "\n",
    "                masked_pose_batch[j][index] = original_pose_batch[select_sample_to_replace][random_replace]\n",
    "                \n",
    "    return masked_pose_batch, masks[:, :, None].repeat(pose_batch.shape[-1], -1), attention_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "cls = torch.zeros(batch_size, 1, 512).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    # reset network\n",
    "    # seed remain the same for the epoch\n",
    "    indexes = np.array(list(range(len(data_train))))\n",
    "    np.random.shuffle(indexes)\n",
    "    _data_train = data_train[indexes]\n",
    "\n",
    "    data_batches = np.array_split(_data_train, len(_data_train)//batch_size)\n",
    "    \n",
    "    # take first batch and prepare\n",
    "    for i, batch in enumerate(data_batches):\n",
    "        pose_batch = norm_train_poses(torch.tensor(np.concatenate([i['pose'][None] for i in batch], 0))).to(device).float()\n",
    "        \n",
    "        masked_pose_batch, mask, attention_mask = generate_masked_batch(pose_batch.clone().detach())\n",
    "        \n",
    "        fake = poseBert(masked_pose_batch, torch.tensor(attention_mask).to(device), device=device, train=True)\n",
    "\n",
    "        errG = criterion(fake.masked_fill(torch.tensor(mask).to(device).bool(), 0), pose_batch.masked_fill(torch.tensor(mask).to(device).bool(), 0))\n",
    "        optimizerG.zero_grad()\n",
    "        errG.backward()\n",
    "        iters += 1\n",
    "        G_losses.append(errG.item())\n",
    "\n",
    "        optimizerG.step()\n",
    "        \n",
    "        if iters % 10 == 0:\n",
    "            print('[%d/%d][%d/%d]\\tLoss_G: %.5f'\n",
    "                  % (epoch, epochs, i, len(data_batches),\n",
    "                     errG.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(poseBert.state_dict(), '/ess_storage/storage/home/skhn2/Poses/models/WORK.MaskedBERT.Ubidata.v2.generator.step{}.G{}.pt'.format(iters, round(errG.item(), 6)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# poseBert.load_state_dict(torch.load('/ess_storage/storage/home/skhn2/Poses/models/WORK.MaskedBERT.Ubidata.generator.step10000.G0.000257.pt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lrG = 1e-4\n",
    "\n",
    "optimizerG = torch.optim.Adam(poseBert.parameters(), lr=lrG, betas=(beta1, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    # reset network\n",
    "    # seed remain the same for the epoch\n",
    "    indexes = np.array(list(range(len(data_train))))\n",
    "    np.random.shuffle(indexes)\n",
    "    _data_train = data_train[indexes]\n",
    "\n",
    "    data_batches = np.array_split(_data_train, len(_data_train)//batch_size)\n",
    "    \n",
    "    # take first batch and prepare\n",
    "    for i, batch in enumerate(data_batches):\n",
    "        pose_batch = norm_train_poses(torch.tensor(np.concatenate([i['pose'][None] for i in batch], 0))).to(device).float()\n",
    "        \n",
    "        masked_pose_batch, mask, attention_mask = generate_masked_batch(pose_batch.clone().detach())\n",
    "        \n",
    "        fake = poseBert(masked_pose_batch, torch.tensor(attention_mask).to(device), device=device, train=True)\n",
    "\n",
    "        errG = criterion(fake.masked_fill(torch.tensor(mask).to(device).bool(), 0), pose_batch.masked_fill(torch.tensor(mask).to(device).bool(), 0))\n",
    "        optimizerG.zero_grad()\n",
    "        errG.backward()\n",
    "        iters += 1\n",
    "        G_losses.append(errG.item())\n",
    "\n",
    "        optimizerG.step()\n",
    "        \n",
    "        if iters % 10 == 0:\n",
    "            print('[%d/%d][%d/%d]\\tLoss_G: %.5f'\n",
    "                  % (epoch, epochs, i, len(data_batches),\n",
    "                     errG.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lrG = 1e-5\n",
    "\n",
    "optimizerG = torch.optim.Adam(poseBert.parameters(), lr=lrG, betas=(beta1, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    # reset network\n",
    "    # seed remain the same for the epoch\n",
    "    indexes = np.array(list(range(len(data_train))))\n",
    "    np.random.shuffle(indexes)\n",
    "    _data_train = data_train[indexes]\n",
    "\n",
    "    data_batches = np.array_split(_data_train, len(_data_train)//batch_size)\n",
    "    \n",
    "    # take first batch and prepare\n",
    "    for i, batch in enumerate(data_batches):\n",
    "        pose_batch = norm_train_poses(torch.tensor(np.concatenate([i['pose'][None] for i in batch], 0))).to(device).float()\n",
    "        \n",
    "        masked_pose_batch, mask, attention_mask = generate_masked_batch(pose_batch.clone().detach())\n",
    "        \n",
    "        fake = poseBert(masked_pose_batch, torch.tensor(attention_mask).to(device), device=device, train=True)\n",
    "\n",
    "        errG = criterion(fake.masked_fill(torch.tensor(mask).to(device).bool(), 0), pose_batch.masked_fill(torch.tensor(mask).to(device).bool(), 0))\n",
    "        optimizerG.zero_grad()\n",
    "        errG.backward()\n",
    "        iters += 1\n",
    "        G_losses.append(errG.item())\n",
    "\n",
    "        optimizerG.step()\n",
    "        \n",
    "        if iters % 10 == 0:\n",
    "            print('[%d/%d][%d/%d]\\tLoss_G: %.5f'\n",
    "                  % (epoch, epochs, i, len(data_batches),\n",
    "                     errG.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    # reset network\n",
    "    # seed remain the same for the epoch\n",
    "    indexes = np.array(list(range(len(data_train))))\n",
    "    np.random.shuffle(indexes)\n",
    "    _data_train = data_train[indexes]\n",
    "\n",
    "    data_batches = np.array_split(_data_train, len(_data_train)//batch_size)\n",
    "    \n",
    "    # take first batch and prepare\n",
    "    for i, batch in enumerate(data_batches):\n",
    "        pose_batch = norm_train_poses(torch.tensor(np.concatenate([i['pose'][None] for i in batch], 0))).to(device).float()\n",
    "        \n",
    "        masked_pose_batch, mask, attention_mask = generate_masked_batch(pose_batch.clone().detach())\n",
    "        \n",
    "        fake = poseBert(masked_pose_batch, torch.tensor(attention_mask).to(device), device=device, train=True)\n",
    "\n",
    "        errG = criterion(fake.masked_fill(torch.tensor(mask).to(device).bool(), 0), pose_batch.masked_fill(torch.tensor(mask).to(device).bool(), 0))\n",
    "        optimizerG.zero_grad()\n",
    "        errG.backward()\n",
    "        iters += 1\n",
    "        G_losses.append(errG.item())\n",
    "\n",
    "        optimizerG.step()\n",
    "        \n",
    "        if iters % 10 == 0:\n",
    "            print('[%d/%d][%d/%d]\\tLoss_G: %.5f'\n",
    "                  % (epoch, epochs, i, len(data_batches),\n",
    "                     errG.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate prediction and save to .bvh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from writers import BVHWriter\n",
    "from pymo.parsers import BVHParser\n",
    "from scipy.spatial.transform import Rotation as R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = BVHParser()\n",
    "# load some existing bvh file\n",
    "placeholder = p.parse('data/placeholders/placeholder.bvh')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# prepare batch from test data\n",
    "pose_batch = norm_train_poses(torch.tensor(np.concatenate([i['pose'][None] for i in [i for i in data if i['type'] == 'test'][:10]], 0))).to(device).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply some masking\n",
    "_masked_pose_batch, _mask, _attention_mask = generate_masked_batch_patched(pose_batch, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    fake = poseBert(_masked_pose_batch.clone().detach()[7:8], torch.tensor(_attention_mask).clone().detach()[7:8].to(device), device=device, train=True)\n",
    "\n",
    "_fake = torch.cat([unnorm_train_poses(pose_batch[:, :100])[7], unnorm_train_poses(fake)[:, :100][0]]).cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# separate global coordinates and motions\n",
    "trace = _fake[:, -3:]\n",
    "pose = _fake[:, :-3].reshape(_fake.shape[0], 22, 3, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert motion from rotation matrix into euler angles\n",
    "pose_euler = np.array([R.from_matrix(i).as_euler('zyx', True) for i in pose]).reshape(pose.shape[0], 66)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, frame in enumerate(pose_euler):\n",
    "    for j, coord in enumerate(frame):\n",
    "        placeholder.values.values[i][j+3] = coord"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, frame in enumerate(trace):\n",
    "    for j, coord in enumerate(frame):\n",
    "        placeholder.values.values[i][j] = coord"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p2 = BVHWriter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('generated/test.bvh', 'w') as f:\n",
    "    p2.write(placeholder, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}