train/train_other.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nWe use FFHQ dataset as training data in this experiment\\n--------------------------------------------------------\\n1) Pre-train the facenet, disentangle nets (part encoders).\\n\\n2) Freeze facenet, disentangle nets (part encoders), part-decoders, and train blending decoders.\\n\\n--------------------------------------------------------\\nAuthor: Peizhi Yan\\nDate Updated (1st): Sep.-17-2021\\nDate Updated (2nd): Oct.-24-2021\\n'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "We use FFHQ dataset as training data in this experiment\n",
    "--------------------------------------------------------\n",
    "1) Pre-train the facenet, disentangle nets (part encoders).\n",
    "\n",
    "2) Freeze facenet, disentangle nets (part encoders), part-decoders, and train blending decoders.\n",
    "\n",
    "--------------------------------------------------------\n",
    "Author: Peizhi Yan\n",
    "Date Updated (1st): Sep.-17-2021\n",
    "Date Updated (2nd): Oct.-24-2021\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\users\\yanpe\\appdata\\local\\programs\\python\\python36\\lib\\site-packages\\numpy\\core\\__init__.py:29: UserWarning: loaded more than 1 DLL from .libs:\n",
      "c:\\users\\yanpe\\appdata\\local\\programs\\python\\python36\\lib\\site-packages\\numpy\\.libs\\libopenblas.TXA6YQSD3GCQQC22GEQ54J2UDCXDXHWN.gfortran-win_amd64.dll\n",
      "c:\\users\\yanpe\\appdata\\local\\programs\\python\\python36\\lib\\site-packages\\numpy\\.libs\\libopenblas.WCDJNK7YVMPZQ2ME2ZZHJJRJ3JIKNDB7.gfortran-win_amd64.dll\n",
      "  stacklevel=1)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA is available. Device:  NVIDIA GeForce RTX 3070\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import gc\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import pickle\n",
    "import scipy.io as sio\n",
    "from scipy.io import loadmat\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import torchvision.transforms as transforms\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "# Pytorch 1.9\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils\n",
    "import torch.distributions\n",
    "\n",
    "# Pytorch3d\n",
    "from pytorch3d.structures import Meshes\n",
    "from pytorch3d.renderer import (\n",
    "    look_at_view_transform,\n",
    "    FoVPerspectiveCameras,\n",
    "    PointLights,\n",
    "    RasterizationSettings,\n",
    "    MeshRenderer,\n",
    "    MeshRasterizer,\n",
    "    SoftPhongShader,\n",
    "    TexturesVertex,\n",
    "    blending\n",
    ")\n",
    "\n",
    "import open3d as o3d\n",
    "\n",
    "# facenet-pytorch 2.5.2\n",
    "from facenet_pytorch import MTCNN, InceptionResnetV1\n",
    "\n",
    "# face-alignment 1.3.4\n",
    "import face_alignment\n",
    "\n",
    "#######################################\n",
    "## Setup PyTorch\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:0\")\n",
    "    torch.cuda.set_device(device)\n",
    "    print('CUDA is available. Device: ', torch.cuda.get_device_name(0))\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print('CUDA is NOT available. Use CPU instead.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10\n"
     ]
    }
   ],
   "source": [
    "img_path = '../datasets/FFHQ/images224x224/{}.png'\n",
    "shape_path = '../datasets/FFHQ/raw_bfm_shape/{}.npy'\n",
    "albedo_path = '../datasets/FFHQ/raw_bfm_color/{}.npy'\n",
    "\n",
    "img_indices = []\n",
    "for fname in os.listdir('../datasets/FFHQ/images224x224/'):\n",
    "    if fname.endswith('.png'):\n",
    "        img_indices.append(fname[:-4])\n",
    "\n",
    "print(len(img_indices))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{1, 2, 3, 4, 5, 6}\n",
      "BFM model loaded\n",
      "\n",
      "S_overall  n_vert:  35709\n",
      "S_eyebrows  n_vert:  444\n",
      "S_eyes  n_vert:  586\n",
      "S_llip  n_vert:  309\n",
      "S_nose  n_vert:  1711\n",
      "S_ulip  n_vert:  576\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    " 1: face skin\n",
    " 2: eye brows\n",
    " 3: eyes\n",
    " 4: nose\n",
    " 5: upper lip\n",
    " 6: lower lip\n",
    "\"\"\"\n",
    "label_map = {\n",
    "    'skin': 1,\n",
    "    'eye_brow': 2,\n",
    "    'eye': 3,\n",
    "    'nose': 4,\n",
    "    'u_lip': 5,\n",
    "    'l_lip': 6\n",
    "}\n",
    "\n",
    "## Load the face parsing labels (per-vertex)\n",
    "vert_labels = np.load('../BFM/bfm_vertex_labels.npy')\n",
    "print(set(vert_labels))\n",
    "\n",
    "\n",
    "## Load the BFM model\n",
    "import pickle\n",
    "with open('../BFM/bfm09.pkl', 'rb') as f:\n",
    "    bfm = pickle.load(f)\n",
    "print('BFM model loaded\\n')\n",
    "\n",
    "## Triangal Facets\n",
    "Faces = bfm['tri'] - 1 ## -1 is critical !!!\n",
    "\n",
    "# find the vertices of part\n",
    "part_vertices = {\n",
    "    'S_overall':[],\n",
    "    'S_eyebrows':[],\n",
    "    'S_eyes':[],\n",
    "    'S_llip':[],\n",
    "    'S_nose':[],\n",
    "    'S_ulip':[]\n",
    "}\n",
    "for idx in range(len(vert_labels)):\n",
    "    part_vertices['S_overall'].append(idx)\n",
    "    if vert_labels[idx] in [label_map['eye_brow']]:\n",
    "        part_vertices['S_eyebrows'].append(idx)\n",
    "    if vert_labels[idx] in [label_map['eye']]:\n",
    "        part_vertices['S_eyes'].append(idx)\n",
    "    if vert_labels[idx] in [label_map['l_lip']]:\n",
    "        part_vertices['S_llip'].append(idx)\n",
    "    if vert_labels[idx] in [label_map['u_lip']]:\n",
    "        part_vertices['S_ulip'].append(idx)\n",
    "    if vert_labels[idx] in [label_map['nose']]:\n",
    "        part_vertices['S_nose'].append(idx)\n",
    "\n",
    "\n",
    "for key in part_vertices:\n",
    "    part_vertices[key] = np.array(part_vertices[key])\n",
    "    print(key, ' n_vert: ', len(part_vertices[key]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'S_overall': (0, 30), 'S_eyebrows': (30, 40), 'S_eyes': (40, 50), 'S_llip': (50, 60), 'S_ulip': (60, 70), 'S_nose': (70, 80)}\n"
     ]
    }
   ],
   "source": [
    "#########################\n",
    "## Latent dimensions\n",
    "latent_dims = {}\n",
    "latent_dims['S_overall'] = 30\n",
    "latent_dims['S_eyebrows'] = 10\n",
    "latent_dims['S_eyes'] = 10\n",
    "latent_dims['S_llip'] = 10\n",
    "latent_dims['S_ulip'] = 10\n",
    "latent_dims['S_nose'] = 10\n",
    "\n",
    "\n",
    "latent_coeff_from_to = {} # later we use the pre-generated indices to merge the part latents \n",
    "_from_ = 0\n",
    "for key in latent_dims:    \n",
    "    latent_coeff_from_to[key] = (_from_, _from_ + latent_dims[key])\n",
    "    _from_ += latent_dims[key]\n",
    "print(latent_coeff_from_to)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decoder S_overall loaded\n",
      "Decoder S_eyebrows loaded\n",
      "Decoder S_eyes loaded\n",
      "Decoder S_llip loaded\n",
      "Decoder S_nose loaded\n",
      "Decoder S_ulip loaded\n",
      "done\n"
     ]
    }
   ],
   "source": [
    "sys.path.append('../')\n",
    "from models.Modules import Encoder, Decoder, VAE, VariationalDisentangleModule,\\\n",
    "OffsetRegressorA, OffsetRegressorB, OffsetRegressor\n",
    "    \n",
    "#################\n",
    "## Part Decoders\n",
    "MODEL_PATH = '../saved_models/part_decoders/{}'\n",
    "part_decoders = {}\n",
    "for key in part_vertices:\n",
    "    part_decoders[key] = Decoder(latent_dim=latent_dims[key], n_vert=len(part_vertices[key])).to(device)\n",
    "    # Load pre-trained parameters\n",
    "    if os.path.exists(MODEL_PATH.format(key)):\n",
    "        part_decoders[key].load_state_dict(torch.load(MODEL_PATH.format(key)))\n",
    "        print('Decoder {} loaded'.format(key))\n",
    "    # freeze the network parameters\n",
    "    for param in part_decoders[key].parameters():\n",
    "        param.requires_grad = False\n",
    "    \n",
    "print('done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model on CUDA:  True\n"
     ]
    }
   ],
   "source": [
    "\n",
    "## Load facenet\n",
    "facenet = InceptionResnetV1(pretrained='vggface2').eval().to(device)\n",
    "for param in facenet.parameters():\n",
    "    param.requires_grad = False\n",
    "    \n",
    "## load our retrained facenet parameters\n",
    "facenet.load_state_dict(torch.load('../saved_models/facenet'))\n",
    "\n",
    "print('model on CUDA: ', next(facenet.parameters()).is_cuda) # True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "S_overall model on CUDA:  True\n",
      "S_eyebrows model on CUDA:  True\n",
      "S_eyes model on CUDA:  True\n",
      "S_llip model on CUDA:  True\n",
      "S_ulip model on CUDA:  True\n",
      "S_nose model on CUDA:  True\n",
      "-----------------\n",
      "S_overall  loaded\n",
      "S_eyebrows  loaded\n",
      "S_eyes  loaded\n",
      "S_llip  loaded\n",
      "S_ulip  loaded\n",
      "S_nose  loaded\n"
     ]
    }
   ],
   "source": [
    "\n",
    "disentangleNets = {}\n",
    "for key in latent_dims:\n",
    "    disentangleNets[key] = VariationalDisentangleModule(512, latent_dims[key]).to(device)\n",
    "    print(key, 'model on CUDA: ', next(disentangleNets[key].parameters()).is_cuda) # True\n",
    "    \n",
    "print('-----------------')\n",
    "    \n",
    "## Load Disentangle Net params\n",
    "disentangle_nets_path = '../saved_models/disentangle_nets/{}'\n",
    "for key in latent_dims:\n",
    "    try:\n",
    "        disentangleNets[key].load_state_dict(torch.load(disentangle_nets_path.format(key)))\n",
    "        print(key, ' loaded')\n",
    "    except:\n",
    "        print(key, ' NOT loaded')\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "##############################\n",
    "## Caching Facenet encodings\n",
    "##############################\n",
    "\n",
    "caching = True\n",
    "\n",
    "if caching:\n",
    "    caching_path = '../cache/facenet_encoding/{}.npy'\n",
    "    batch_size = 16\n",
    "    facenet.eval()\n",
    "    batch_x = torch.zeros([batch_size, 3, 224, 224], dtype=torch.float32).to(device)\n",
    "\n",
    "    i = 0\n",
    "    with tqdm(total=len(img_indices)) as pbar:\n",
    "        while i < len(img_indices):\n",
    "            f_indices = []\n",
    "            j = 0\n",
    "            while j < batch_size and i < len(img_indices):            \n",
    "                file_index = img_indices[i]\n",
    "                #if os.path.exists(caching_path.format(file_index)):\n",
    "                #    i += 1\n",
    "                #    continue\n",
    "                try:\n",
    "                    img = Image.open(img_path.format(file_index))\n",
    "                    img = img.resize((224,224))\n",
    "                    img = np.asarray(img, dtype=np.float32)\n",
    "                    img = torch.from_numpy(img)\n",
    "                    img = img.permute(2,0,1)\n",
    "                except:\n",
    "                    i += 1\n",
    "                    continue\n",
    "                batch_x[j] = img\n",
    "                f_indices.append(file_index)\n",
    "                i += 1\n",
    "                j += 1\n",
    "            if len(f_indices) == 0:\n",
    "                continue\n",
    "            img_encodings = facenet(batch_x).detach().cpu().numpy()\n",
    "            for j in range(len(f_indices)):\n",
    "                file_index = f_indices[j]\n",
    "                np.save(caching_path.format(file_index), img_encodings[j])\n",
    "            pbar.update(batch_size)\n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def next_batch(idx, batch_size, img_path, shape_path, albedo_path, img_indices):\n",
    "    \"\"\"\n",
    "        return: images, shapes, albedos, new index\n",
    "    \"\"\"\n",
    "    if idx + batch_size > len(img_indices):\n",
    "        batch_size = len(img_indices) - idx\n",
    "                \n",
    "    batch_x = torch.zeros([batch_size, 3, 224, 224], dtype=torch.float32).to(device)\n",
    "    batch_s = torch.zeros([batch_size, 35709, 3], dtype=torch.float32).to(device)\n",
    "    batch_t = torch.zeros([batch_size, 35709, 3], dtype=torch.float32).to(device)\n",
    "\n",
    "    i = idx\n",
    "    counter = 0\n",
    "    while i < idx + batch_size:\n",
    "        file_index = img_indices[i]\n",
    "        \n",
    "        img = Image.open(img_path.format(file_index))\n",
    "        img = img.resize((224,224))\n",
    "        img = np.asarray(img, dtype=np.float32)\n",
    "        img = torch.from_numpy(img)\n",
    "        img = img.permute(2,0,1)\n",
    "        \n",
    "        batch_x[counter] = img\n",
    "        batch_s[counter] = torch.from_numpy(np.reshape(np.load(shape_path.format(file_index)), [35709, 3]))\n",
    "        batch_t[counter] = torch.from_numpy(np.reshape(np.load(albedo_path.format(file_index))/255., [35709, 3]))\n",
    "        \n",
    "        i += 1\n",
    "        counter += 1\n",
    "        \n",
    "    return batch_x, batch_s, batch_t, i\n",
    "\n",
    "\n",
    "def next_batch_cached(idx, batch_size, encoding_path, shape_path, albedo_path, img_indices):\n",
    "    \"\"\" \n",
    "    Use cached facenet encodings \n",
    "        return: img encodings, shapes, albedos, new index\n",
    "    \"\"\"\n",
    "    if idx + batch_size > len(img_indices):\n",
    "        batch_size = len(img_indices) - idx\n",
    "                \n",
    "    batch_e = torch.zeros([batch_size, 512], dtype=torch.float32).to(device)\n",
    "    batch_s = torch.zeros([batch_size, 35709, 3], dtype=torch.float32).to(device)\n",
    "    batch_t = torch.zeros([batch_size, 35709, 3], dtype=torch.float32).to(device)\n",
    "\n",
    "    i = idx\n",
    "    counter = 0\n",
    "    while i < idx + batch_size:\n",
    "        file_index = img_indices[i]\n",
    "        \n",
    "        batch_e[counter] = torch.from_numpy(np.reshape(np.load(encoding_path.format(file_index)), [512]))\n",
    "        batch_s[counter] = torch.from_numpy(np.reshape(np.load(shape_path.format(file_index)), [35709, 3]))\n",
    "        batch_t[counter] = torch.from_numpy(np.reshape(np.load(albedo_path.format(file_index))/255., [35709, 3]))\n",
    "        \n",
    "        i += 1\n",
    "        counter += 1\n",
    "        \n",
    "    return batch_e, batch_s, batch_t, i\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_shape_loss(pred_shape, targ_shape):\n",
    "    batch_size = pred_shape.shape[0]\n",
    "    loss = torch.square(pred_shape - targ_shape)\n",
    "    loss = loss.sum(dim=[0, 1, 2]) / batch_size\n",
    "    #loss = torch.mean(loss)\n",
    "    return loss\n",
    "\n",
    "def standardize_part_shape(part_shapes):\n",
    "    # part_shapes: [n, n_vert, 3]  n -- batch size\n",
    "    y_max, _ = torch.max(part_shapes[...,1], dim=1) # [n, 1]\n",
    "    y_min, _ = torch.min(part_shapes[...,1], dim=1) # [n, 1]\n",
    "    y_center = (y_max + y_min) / 2 # [n, 1]\n",
    "    z_max, _ = torch.max(part_shapes[...,2], dim=1) # [n, 1]\n",
    "    z_min, _ = torch.min(part_shapes[...,2], dim=1) # [n, 1]\n",
    "    z_center = (z_max + z_min) / 2 # [n, 1]\n",
    "    batch_size = part_shapes.shape[0]\n",
    "    for i in range(batch_size):\n",
    "        part_shapes[i, :, 1] -= y_center[i]\n",
    "        part_shapes[i, :, 2] -= z_center[i]\n",
    "    return part_shapes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "tb_writer = SummaryWriter()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train disentangle layers (part encoders)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#######################################\n",
    "## Disentangle Layers Pre-Training\n",
    "#######################################\n",
    "\n",
    "EPOCHS = 100\n",
    "#BATCH_SIZE = 128\n",
    "BATCH_SIZE = 8\n",
    "\n",
    "caching_path = './cache/facenet_encoding/{}.npy'\n",
    "\n",
    "\n",
    "####################\n",
    "## Some switches\n",
    "SW = {0: False, 1: True} # Switch\n",
    "train_or_not = {} # indicate whether or not to train a part module\n",
    "train_or_not['Facenet'] =      SW[0]\n",
    "use_cached_facenet_encodings = SW[1]\n",
    "train_or_not['disentangle']  = SW[1]\n",
    "train_or_not['part_decoders'] = SW[0] # please do not train part-decoders!!\n",
    "train_or_not['S_overall'] =     SW[1]\n",
    "train_or_not['S_eyebrows'] = SW[1]\n",
    "train_or_not['S_eyes'] =     SW[1]\n",
    "train_or_not['S_llip'] =     SW[1]\n",
    "train_or_not['S_ulip'] =     SW[1]\n",
    "train_or_not['S_nose'] =     SW[1]\n",
    "\n",
    "parameters = [] ## the parameters to be optimized\n",
    "\n",
    "################################\n",
    "## Disentangle Layer parameters\n",
    "for key in latent_dims:\n",
    "    if train_or_not['disentangle']:\n",
    "        if train_or_not[key] == False:\n",
    "            continue\n",
    "        if parameters is None:\n",
    "            parameters = list(disentangleNets[key].parameters())\n",
    "        else:\n",
    "            parameters += list(disentangleNets[key].parameters())\n",
    "\n",
    "################################\n",
    "## Unfreeze FaceNet parameters\n",
    "if train_or_not['Facenet']:\n",
    "    facenet.train()\n",
    "    for param in facenet.parameters():\n",
    "        param.requires_grad = True\n",
    "    parameters += list(facenet.parameters())   \n",
    "else:\n",
    "    facenet.eval()\n",
    "    for param in facenet.parameters():\n",
    "        param.requires_grad = False\n",
    "\n",
    "################################\n",
    "## Part Decoders\n",
    "if train_or_not['part_decoders']:\n",
    "    for key in latent_dims:\n",
    "        if train_or_not[key] == False:\n",
    "            continue\n",
    "        for param in part_decoders[key].parameters():\n",
    "            param.requires_grad = True\n",
    "        parameters += list(part_decoders[key].parameters())    \n",
    "\n",
    "###############\n",
    "## Optimizers\n",
    "lr = 1e-4 # initial learning rate\n",
    "#lr = 1e-5\n",
    "optimizers = {}\n",
    "for key in latent_dims:\n",
    "    disentangle = disentangleNets[key]\n",
    "    optimizers[key] = torch.optim.Adam(disentangle.parameters(), lr=lr)\n",
    "    \n",
    "#####################\n",
    "## Gaussian Sampler\n",
    "normal = torch.distributions.Normal(0, 1) # a normal distribution with mean = 0, std = 1\n",
    "if device.type == 'cuda':\n",
    "    normal.loc = normal.loc.cuda() # use CUDA GPU for sampling\n",
    "    normal.scale = normal.scale.cuda()\n",
    "\n",
    "\n",
    "#############\n",
    "## Training\n",
    "iter_ = 0\n",
    "for epoch in range(EPOCHS):\n",
    "    \n",
    "    idx = 100 # first 100 images as validation data\n",
    "    with tqdm(total=len(img_indices[100:])) as pbar:\n",
    "        while idx < len(img_indices[100:]):\n",
    "\n",
    "\n",
    "            ######################\n",
    "            # Prepare batch data\n",
    "            if use_cached_facenet_encodings:\n",
    "                img_encodings, batch_s, batch_t, idx = next_batch_cached(idx, BATCH_SIZE, caching_path, \n",
    "                                                                   shape_path, albedo_path, img_indices)                \n",
    "            else:\n",
    "                batch_x, batch_s, batch_t, idx = next_batch(idx, BATCH_SIZE, img_path, shape_path, albedo_path, img_indices)\n",
    "                if batch_x is None:\n",
    "                    continue\n",
    "                img_encodings = facenet(batch_x)\n",
    "                \n",
    "            ##################\n",
    "            # Train each part\n",
    "            for key in latent_dims:\n",
    "                if train_or_not[key] == False:\n",
    "                    continue\n",
    "\n",
    "                ##################\n",
    "                # Clean gradient\n",
    "                optimizers[key].zero_grad()\n",
    "                    \n",
    "                part_mu, part_sigma = disentangleNets[key](img_encodings)\n",
    "                part_latents = part_mu + part_sigma * normal.sample(part_mu.size())\n",
    "                kl_loss = torch.mean(torch.sum(-0.5 * (1 + part_sigma - part_mu**2 - torch.exp(part_sigma)), dim=1))\n",
    "                preds = part_decoders[key](part_latents)\n",
    "\n",
    "                batch_shape_part = batch_s[:,part_vertices[key],:]\n",
    "                if key != 'S_overall':\n",
    "                    ## Standardize the part shapes\n",
    "                    batch_shape_part = standardize_part_shape(batch_shape_part)\n",
    "                part_shape_loss = compute_shape_loss(preds, batch_shape_part)\n",
    "                    \n",
    "                loss = part_shape_loss + 1e3*kl_loss\n",
    "                                            \n",
    "                ######\n",
    "                # BP\n",
    "                loss.backward()\n",
    "\n",
    "                ############\n",
    "                # Optimize\n",
    "                optimizers[key].step()\n",
    "                \n",
    "                ################\n",
    "                ## TensorBoard\n",
    "                if iter_ % 100 == 0:\n",
    "                    tb_writer.add_scalar('{}/res_loss'.format(key), part_shape_loss.data.detach().cpu().numpy(), iter_)\n",
    "                    tb_writer.add_scalar('{}/kl_loss'.format(key), kl_loss.data.detach().cpu().numpy(), iter_)\n",
    "                    tb_writer.add_scalar('{}/loss'.format(key), loss.data.detach().cpu().numpy(), iter_)\n",
    "\n",
    "            \n",
    "            \n",
    "            pbar.update(BATCH_SIZE)\n",
    "            iter_ += 1\n",
    "\n",
    "    print('Epoch: ', epoch)\n",
    "\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train facenet + disentangle layers (part encoders)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#######################################\n",
    "## Fine-tuning\n",
    "#######################################\n",
    "\n",
    "EPOCHS = 100\n",
    "#BATCH_SIZE = 128\n",
    "BATCH_SIZE = 8\n",
    "\n",
    "caching_path = './cache/facenet_encoding/{}.npy'\n",
    "\n",
    "\n",
    "####################\n",
    "## Some switches\n",
    "SW = {0: False, 1: True} # Switch\n",
    "train_or_not = {} # indicate whether or not to train a part module\n",
    "train_or_not['Facenet'] =      SW[1]\n",
    "use_cached_facenet_encodings = SW[0]\n",
    "train_or_not['disentangle']  = SW[1]\n",
    "train_or_not['part_decoders'] = SW[0] # please do not train part-decoders!!\n",
    "train_or_not['S_overall'] =     SW[1]\n",
    "train_or_not['S_eyebrows'] = SW[1]\n",
    "train_or_not['S_eyes'] =     SW[1]\n",
    "train_or_not['S_llip'] =     SW[1]\n",
    "train_or_not['S_ulip'] =     SW[1]\n",
    "train_or_not['S_nose'] =     SW[1]\n",
    "\n",
    "parameters = [] ## the parameters to be optimized\n",
    "\n",
    "################################\n",
    "## Disentangle Layer parameters\n",
    "for key in latent_dims:\n",
    "    if train_or_not['disentangle']:\n",
    "        if train_or_not[key] == False:\n",
    "            continue\n",
    "        if parameters is None:\n",
    "            parameters = list(disentangleNets[key].parameters())\n",
    "        else:\n",
    "            parameters += list(disentangleNets[key].parameters())\n",
    "\n",
    "################################\n",
    "## Unfreeze FaceNet parameters\n",
    "if train_or_not['Facenet']:\n",
    "    facenet.train()\n",
    "    for param in facenet.parameters():\n",
    "        param.requires_grad = True\n",
    "    parameters += list(facenet.parameters())   \n",
    "else:\n",
    "    facenet.eval()\n",
    "    for param in facenet.parameters():\n",
    "        param.requires_grad = False\n",
    "\n",
    "################################\n",
    "## Part Decoders\n",
    "if train_or_not['part_decoders']:\n",
    "    for key in latent_dims:\n",
    "        if train_or_not[key] == False:\n",
    "            continue\n",
    "        for param in part_decoders[key].parameters():\n",
    "            param.requires_grad = True\n",
    "        parameters += list(part_decoders[key].parameters())    \n",
    "\n",
    "###############\n",
    "## Optimizers\n",
    "#lr = 1e-4 # initial learning rate\n",
    "lr = 1e-5\n",
    "optimizers = {}\n",
    "for key in latent_dims:\n",
    "    disentangle = disentangleNets[key]\n",
    "    optimizers[key] = torch.optim.Adam(disentangle.parameters(), lr=lr)\n",
    "    \n",
    "#####################\n",
    "## Gaussian Sampler\n",
    "normal = torch.distributions.Normal(0, 1) # a normal distribution with mean = 0, std = 1\n",
    "if device.type == 'cuda':\n",
    "    normal.loc = normal.loc.cuda() # use CUDA GPU for sampling\n",
    "    normal.scale = normal.scale.cuda()\n",
    "\n",
    "\n",
    "#############\n",
    "## Training\n",
    "iter_ = 0\n",
    "for epoch in range(EPOCHS):\n",
    "    \n",
    "    idx = 100 # first 100 images as validation data\n",
    "    with tqdm(total=len(img_indices[100:])) as pbar:\n",
    "        while idx < len(img_indices[100:]):\n",
    "\n",
    "\n",
    "            ######################\n",
    "            # Prepare batch data\n",
    "            if use_cached_facenet_encodings:\n",
    "                img_encodings, batch_s, batch_t, idx = next_batch_cached(idx, BATCH_SIZE, caching_path, \n",
    "                                                                   shape_path, albedo_path, img_indices)                \n",
    "            else:\n",
    "                batch_x, batch_s, batch_t, idx = next_batch(idx, BATCH_SIZE, img_path, shape_path, albedo_path, img_indices)\n",
    "                if batch_x is None:\n",
    "                    continue\n",
    "                img_encodings = facenet(batch_x)\n",
    "                \n",
    "            ##################\n",
    "            # Train each part\n",
    "            for key in latent_dims:\n",
    "                if train_or_not[key] == False:\n",
    "                    continue\n",
    "\n",
    "                ##################\n",
    "                # Clean gradient\n",
    "                optimizers[key].zero_grad()\n",
    "                    \n",
    "                part_mu, part_sigma = disentangleNets[key](img_encodings)\n",
    "                part_latents = part_mu + part_sigma * normal.sample(part_mu.size())\n",
    "                kl_loss = torch.mean(torch.sum(-0.5 * (1 + part_sigma - part_mu**2 - torch.exp(part_sigma)), dim=1))\n",
    "                preds = part_decoders[key](part_latents)\n",
    "\n",
    "                batch_shape_part = batch_s[:,part_vertices[key],:]\n",
    "                if key != 'S_overall':\n",
    "                    ## Standardize the part shapes\n",
    "                    batch_shape_part = standardize_part_shape(batch_shape_part)\n",
    "                part_shape_loss = compute_shape_loss(preds, batch_shape_part)\n",
    "                    \n",
    "                loss = part_shape_loss + 1e3*kl_loss\n",
    "                                            \n",
    "                ######\n",
    "                # BP\n",
    "                loss.backward()\n",
    "\n",
    "                ############\n",
    "                # Optimize\n",
    "                optimizers[key].step()\n",
    "                \n",
    "                ################\n",
    "                ## TensorBoard\n",
    "                if iter_ % 100 == 0:\n",
    "                    tb_writer.add_scalar('{}/res_loss'.format(key), part_shape_loss.data.detach().cpu().numpy(), iter_)\n",
    "                    tb_writer.add_scalar('{}/kl_loss'.format(key), kl_loss.data.detach().cpu().numpy(), iter_)\n",
    "                    tb_writer.add_scalar('{}/loss'.format(key), loss.data.detach().cpu().numpy(), iter_)\n",
    "\n",
    "            \n",
    "            \n",
    "            pbar.update(BATCH_SIZE)\n",
    "            iter_ += 1\n",
    "\n",
    "    print('Epoch: ', epoch)\n",
    "\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Save all the networks \n",
    "\n",
    "SAVE = False\n",
    "\n",
    "if SAVE:\n",
    "    model_save_path = '../saved_models/disentangle_nets/{}'\n",
    "    for key in latent_dims:\n",
    "        torch.save(disentangleNets[key].state_dict(), model_save_path.format(key))\n",
    "        print(key, ' saved')\n",
    "        \n",
    "    model_save_path = '../saved_models/facenet'\n",
    "    for key in latent_dims:\n",
    "        torch.save(facenet.state_dict(), model_save_path.format(key))\n",
    "        print('facenet saved')\n",
    "\n",
    "\n",
    "\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train offset regressor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##############################\n",
    "## Offset Prediction Module\n",
    "##############################\n",
    "offsetKeys = {'S_eyebrows': 0, \n",
    "              'S_eyes': 1, \n",
    "              'S_llip': 2, \n",
    "              'S_ulip': 3, \n",
    "              'S_nose': 4}\n",
    "\n",
    "offsetRegressorA = OffsetRegressorA().to(device)\n",
    "\n",
    "offserRegressorsB = {}\n",
    "for key in offsetKeys:\n",
    "    offserRegressorsB[key+'-y'] = OffsetRegressorB()\n",
    "    offserRegressorsB[key+'-z'] = OffsetRegressorB()\n",
    "\n",
    "offsetRegressor = OffsetRegressor(offsetRegressorA, offserRegressorsB).to(device)\n",
    "\n",
    "## Load the pre-trained decoder weights\n",
    "#offsetRegressor.load_state_dict(torch.load('../saved_models/offset_regressor'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = [] ## the parameters to be optimized\n",
    "\n",
    "\n",
    "################################\n",
    "## Offset Decoders\n",
    "for param in offsetRegressor.parameters():\n",
    "    param.requires_grad = True\n",
    "    \n",
    "#parameters += list(offsetRegressor.parameters())\n",
    "parameters += list(offsetRegressor.S_eyes_y.parameters())\n",
    "parameters += list(offsetRegressor.S_eyes_z.parameters())\n",
    "\n",
    "    \n",
    "###############\n",
    "## Optimizer\n",
    "optimizer = torch.optim.Adam(parameters, lr=1e-4) # initial learning rate\n",
    "#optimizer = torch.optim.Adam(parameters, lr=1e-5) # \n",
    "\n",
    "\n",
    "#############\n",
    "## Training\n",
    "for epoch in range(EPOCHS):\n",
    "    \n",
    "    idx = 100 # first 100 images as validation data\n",
    "    with tqdm(total=len(img_indices[100:])) as pbar:\n",
    "        while idx < len(img_indices[100:]):\n",
    "            ##################\n",
    "            # Clean gradient\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            ######################\n",
    "            # Prepare batch data\n",
    "            img_encodings, batch_s, idx = next_batch_cached(idx, BATCH_SIZE, caching_path, \n",
    "                                                                     shape_path, img_indices)                \n",
    "\n",
    "                \n",
    "            ###########\n",
    "            # Predict\n",
    "            pred_offsets = offsetRegressor(img_encodings)\n",
    "            \n",
    "                \n",
    "            ################\n",
    "            # Compute loss\n",
    "            batch_offsets = torch.zeros([BATCH_SIZE, 5, 2], dtype=torch.float32).to(device)\n",
    "            for key in latent_dims:\n",
    "                if key[0] == 'S':\n",
    "                    if key != 'S_rest':\n",
    "                        batch_shape_part = batch_s[:,part_vertices[key],:]\n",
    "                        ## Standardize the part shapes\n",
    "                        _, batch_y_offset, batch_z_offset = standardize_part_shape(batch_shape_part)\n",
    "                        batch_offsets[:, offsetKeys[key], 0] = batch_y_offset\n",
    "                        batch_offsets[:, offsetKeys[key], 1] = batch_z_offset            \n",
    "            offset_loss = torch.sum(torch.square(pred_offsets - batch_offsets), dim=[0,1,2]) / BATCH_SIZE\n",
    "            loss = offset_loss\n",
    "                \n",
    "            ######\n",
    "            # BP\n",
    "            loss.backward()\n",
    "\n",
    "            ############\n",
    "            # Optimize\n",
    "            optimizer.step()\n",
    "            \n",
    "            pbar.update(BATCH_SIZE)\n",
    "\n",
    "    print('Epoch: ', epoch)\n",
    "    print(' -- Overall loss: ' , loss.data)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Save offser regressor network\n",
    "\n",
    "SAVE = False\n",
    "\n",
    "if SAVE:\n",
    "    model_save_path = '../saved_models/offset_regressor'\n",
    "    torch.save(offsetRegressor.state_dict(), model_save_path)\n",
    "    print('offset regressor saved')\n",
    "            \n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "Faces = bfm['tri'] - 1 ## -1 is critical !!!\n",
    "\n",
    "def render(V, T, Faces, width=512, height=512):\n",
    "    ###############################\n",
    "    ## Visualize the render result\n",
    "    o3d_mesh = o3d.geometry.TriangleMesh()\n",
    "    o3d_mesh.vertices = o3d.utility.Vector3dVector(V) # dtype vector3d (float)\n",
    "    o3d_mesh.triangles = o3d.utility.Vector3iVector(Faces) # dtype vector3i (int)\n",
    "    if T is not None:\n",
    "        o3d_mesh.vertex_colors = o3d.utility.Vector3dVector(T) # dtype vector3i (int)\n",
    "    o3d_mesh.compute_vertex_normals() # computing normal will give specular effect while rendering\n",
    "\n",
    "    vis = o3d.visualization.Visualizer()\n",
    "    vis.create_window(width=width, height=height, visible = False)\n",
    "    vis.add_geometry(o3d_mesh)\n",
    "    #depth = vis.capture_depth_float_buffer(True)\n",
    "    image = vis.capture_screen_float_buffer(True)\n",
    "\n",
    "    return o3d_mesh, image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "############################\n",
    "## Reconstruction Testing\n",
    "############################\n",
    "\n",
    "facenet.eval()\n",
    "for param in facenet.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "idx = 0\n",
    "#len(img_indices)\n",
    "for idx in range(100):\n",
    "    #idx += 1\n",
    "\n",
    "    # Prepare batch data\n",
    "    batch_x, batch_s, batch_t, _ = next_batch(idx, 1, img_path, shape_path, albedo_path, img_indices)\n",
    "    if batch_x is None:\n",
    "        continue\n",
    "    \n",
    "    # Predict\n",
    "    preds = {}\n",
    "    img_encodings = facenet(batch_x)\n",
    "    for key in latent_dims:\n",
    "        part_mu, part_sigma = disentangleNets[key](img_encodings)\n",
    "        part_latents = part_mu\n",
    "        preds[key] = part_decoders[key](part_latents).detach().cpu().numpy()\n",
    "\n",
    "    pred_offsets = offsetRegressor(img_encodings).detach().cpu().numpy()\n",
    "\n",
    "    for key in offsetKeys:\n",
    "        preds[key][0,:,1:] += pred_offsets[0, offsetKeys[key], :]\n",
    "        \n",
    "        \n",
    "    V_targ = np.zeros([35709, 3])\n",
    "    V_targ = batch_s.detach().cpu().numpy()[0]\n",
    "    T_targ = np.zeros([35709, 3])\n",
    "    T_targ = batch_t.detach().cpu().numpy()[0]\n",
    "\n",
    "    V_pred = np.zeros([35709, 3])\n",
    "    for key in latent_dims:\n",
    "        V_pred[part_vertices[key]] = preds[key][0]\n",
    "\n",
    "\n",
    "    plt.figure(figsize=(10,5))\n",
    "    o3d_mesh_targ, image_targ = render(V_targ, None, Faces)\n",
    "    plt.subplot(1,4,1)\n",
    "    plt.title('Target Mesh')\n",
    "    plt.imshow(np.asarray(image_targ))\n",
    "\n",
    "    o3d_mesh_pred, image_pred = render(V_pred, None, Faces)\n",
    "    plt.subplot(1,4,2)\n",
    "    plt.title('Predicted Mesh')\n",
    "    plt.imshow(np.asarray(image_pred))\n",
    "    \n",
    "    \"\"\"\n",
    "    o3d_mesh_pred, image_pred = render(V_pred, T_pred, Faces)\n",
    "    plt.subplot(1,4,3)\n",
    "    plt.title('Predicted Mesh w/ Color')\n",
    "    plt.imshow(np.asarray(image_pred))\n",
    "    \"\"\"\n",
    "    \n",
    "    plt.subplot(1,4,4)\n",
    "    plt.title('{}'.format(idx))\n",
    "    plt.imshow(batch_x[0].permute(1,2,0).detach().cpu().numpy()/255)\n",
    "    plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "############################\n",
    "## Latent Swapping Demo\n",
    "############################\n",
    "\n",
    "\n",
    "def demo(idx_A, idx_B, result_combination):\n",
    "    batch_x_A, batch_s_A, batch_t_A, _ = next_batch(idx_A, 1, img_path, shape_path, albedo_path, img_indices)\n",
    "    batch_x_B, batch_s_B, batch_t_B, _ = next_batch(idx_B, 1, img_path, shape_path, albedo_path, img_indices)\n",
    "\n",
    "    shape_latents_A = torch.zeros([1, shape_latent_dim], dtype=torch.float32).to(device) # as input of shape blending decoder\n",
    "    shape_latents_B = torch.zeros([1, shape_latent_dim], dtype=torch.float32).to(device) # as input of shape blending decoder\n",
    "    shape_latents_C = torch.zeros([1, shape_latent_dim], dtype=torch.float32).to(device) # as input of shape blending decoder\n",
    "    preds_A = {}\n",
    "    preds_B = {}\n",
    "    preds_C = {}\n",
    "    \n",
    "    pointer = {\n",
    "        'A': [batch_x_A, batch_s_A, batch_t_A, shape_latents_A, preds_A],\n",
    "        'B': [batch_x_B, batch_s_B, batch_t_B, shape_latents_B, preds_B],\n",
    "        'C': [None, None, None, shape_latents_C, preds_C]\n",
    "    }\n",
    "\n",
    "    for pnt in pointer:\n",
    "        if pnt == 'C':\n",
    "            continue\n",
    "        [batch_x, batch_s, batch_t, shape_latents, preds] = pointer[pnt]\n",
    "        for key in latent_dims:\n",
    "            img_encodings = facenet(batch_x)\n",
    "            part_mu, part_sigma = disentangleNets[key](img_encodings)\n",
    "            part_latents = part_mu\n",
    "            preds[key] = part_decoders[key](part_latents).detach().cpu().numpy()\n",
    "            (_from_, _to_) = latent_coeff_from_to[key]\n",
    "            if key[0] == 'S':\n",
    "                shape_latents[:, _from_:_to_] = part_latents\n",
    "\n",
    "    for key in latent_dims:\n",
    "        source = result_combination[key]\n",
    "        [_, _, _, shape_latents_source, preds_source] = pointer[source]\n",
    "        (_from_, _to_) = latent_coeff_from_to[key]    \n",
    "        if key[0] == 'S':\n",
    "            shape_latents_C[:, _from_:_to_] = shape_latents_source[:, _from_:_to_]\n",
    "        preds_C[key] = preds_source[key]\n",
    "        \n",
    "    for pnt in pointer:\n",
    "        [_, _, _, shape_latents, preds] = pointer[pnt]\n",
    "        pred_offsets = offsetDecoder(shape_latents).detach().cpu().numpy()\n",
    "        preds['offsets'] = pred_offsets\n",
    "        #for key in offsetKeys:\n",
    "        #    preds[key][0,:,1:] += pred_offsets[0, offsetKeys[key], :]\n",
    "\n",
    "\n",
    "    ## Visualize results\n",
    "    plt.figure(figsize=(15,5))\n",
    "    i = 1\n",
    "    meshes = []\n",
    "    for pnt in pointer:\n",
    "        [_, _, _, shape_latents, preds] = pointer[pnt]\n",
    "        V_pred = np.zeros([35709, 3], dtype=np.float32)\n",
    "        T_pred = np.zeros([35709, 3], dtype=np.float32)\n",
    "        for key in latent_dims:\n",
    "            if key[0] == 'S':\n",
    "                pred_part = np.copy(preds[key])\n",
    "                if key in offsetKeys:\n",
    "                    pred_part[0,:,1:] += preds['offsets'][0, offsetKeys[key], :]\n",
    "                V_pred[part_vertices[key]] = pred_part\n",
    "            else:\n",
    "                T_pred[part_vertices[key]] = preds[key]\n",
    "        pred_mesh, rendered_img = render(V_pred, T_pred, Faces)\n",
    "        meshes.append(pred_mesh)\n",
    "        plt.subplot(1,3,i); i+=1\n",
    "        plt.title('{}'.format(pnt))\n",
    "        plt.imshow(np.asarray(rendered_img))\n",
    "    \n",
    "    return meshes[0], meshes[1], meshes[2]\n",
    "    \n",
    "    \n",
    "idx_A = 20\n",
    "idx_B = 29\n",
    "result_combination = {} \n",
    "result_combination['S_rest'] = 'A'\n",
    "result_combination['S_eyebrows'] = 'A'\n",
    "result_combination['S_eyes'] = 'B'\n",
    "result_combination['S_llip'] = 'B'\n",
    "result_combination['S_ulip'] = 'B'\n",
    "result_combination['S_nose'] = 'B'\n",
    "result_combination['T_eyebrows'] = 'A' \n",
    "result_combination['T_eyes'] = 'B'\n",
    "result_combination['T_lips'] = 'B'\n",
    "result_combination['T_skin'] = 'A'\n",
    "\n",
    "mesh_A, mesh_B, mesh_C = demo(idx_A, idx_B, result_combination)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "o3d.visualization.draw_geometries([mesh_C]) \n"
   ]
  }
 ],
 "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.7rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}