Live_Speech_Portraits.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Live Speech Portraits",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tg-bomze/collection-of-notebooks/blob/master/Live_Speech_Portraits.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BD6QsxyHlu4z"
      },
      "source": [
        "<b><font color=\"black\" size=\"+4\">LiveSpeechPortraits</font></b>\n",
        "\n",
        "<b><font color=\"black\" size=\"+2\">Based on:</font></b>\n",
        "\n",
        "**GitHub repository**: [LiveSpeechPortraits](https://github.com/YuanxunLu/LiveSpeechPortraits)\n",
        "\n",
        "Creator: **[OldSix](https://github.com/YuanxunLu)**\n",
        "\n",
        "Article: [Live Speech Portraits: Real-Time Photorealistic Talking-Head Animation](https://arxiv.org/pdf/2109.10595.pdf).\n",
        "\n",
        "<b><font color=\"black\" size=\"+2\">Colab created by:</font></b>\n",
        "\n",
        "GitHub: [@tg-bomze](https://github.com/tg-bomze),\n",
        "Telegram: [@bomze](https://t.me/bomze),\n",
        "Twitter: [@tg_bomze](https://twitter.com/tg_bomze).\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "```\n",
        "(ENG) To get started, click on the button (where the red arrow indicates). After clicking, wait until the execution is complete.\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K3EWp0HtQea1"
      },
      "source": [
        "<b><font color=\"red\" size=\"+3\">Attention</font></b>\n",
        "\n",
        "**Follow link**: https://drive.google.com/drive/folders/1sHc2xEEGwnb0h2rkUhG9sPmOxvRvPVpJ - add shortcut to Drive (**'data' folder**)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8VrBxQ9gQwFc",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Install</font></b>\n",
        "from google.colab import drive\n",
        "from IPython.display import clear_output\n",
        "from google.colab import files\n",
        "import os\n",
        "\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "!git clone https://github.com/YuanxunLu/LiveSpeechPortraits.git\n",
        "%cd /content/LiveSpeechPortraits\n",
        "!pip install -r requirements.txt\n",
        "!pip install dominate\n",
        "!pip install sox\n",
        "\n",
        "!rm -rf /content/LiveSpeechPortraits/data\n",
        "!cp -r /content/drive/MyDrive/data /content/LiveSpeechPortraits\n",
        "clear_output()"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "0jVzcxc4e53D"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Fix demo.py</font></b>\n",
        "%%writefile /content/LiveSpeechPortraits/demo.py\n",
        "import os\n",
        "import subprocess\n",
        "from os.path import join\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import torch\n",
        "from collections import OrderedDict\n",
        "import librosa\n",
        "from skimage.io import imread\n",
        "import cv2\n",
        "import scipy.io as sio\n",
        "import argparse\n",
        "import yaml\n",
        "import albumentations as A\n",
        "import albumentations.pytorch\n",
        "from pathlib import Path\n",
        "\n",
        "from options.test_audio2feature_options import TestOptions as FeatureOptions\n",
        "from options.test_audio2headpose_options import TestOptions as HeadposeOptions\n",
        "from options.test_feature2face_options import TestOptions as RenderOptions\n",
        "\n",
        "from datasets import create_dataset\n",
        "from models import create_model\n",
        "from models.networks import APC_encoder\n",
        "import util.util as util\n",
        "from util.visualizer import Visualizer\n",
        "from funcs import utils\n",
        "from funcs import audio_funcs\n",
        "import soundfile as sf\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "\n",
        "\n",
        "def write_video_with_audio(audio_path, output_path, prefix='pred_'):\n",
        "    fps, fourcc = 60, cv2.VideoWriter_fourcc(*'DIVX')\n",
        "    video_tmp_path = join(save_root, 'tmp.avi')\n",
        "    out = cv2.VideoWriter(video_tmp_path, fourcc, fps, (Renderopt.loadSize, Renderopt.loadSize))\n",
        "    for j in tqdm(range(nframe), position=0, desc='writing video'):\n",
        "        img = cv2.imread(join(save_root, prefix + str(j+1) + '.jpg'))\n",
        "        out.write(img)\n",
        "    out.release() \n",
        "    cmd = 'ffmpeg -i \"' + video_tmp_path + '\" -i \"' + audio_path + '\" -codec copy -shortest \"' + output_path + '\"'\n",
        "    subprocess.call(cmd, shell=True) \n",
        "    os.remove(video_tmp_path)  # remove the template video\n",
        "\n",
        "\n",
        "\n",
        "if __name__ == '__main__': \n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument('--id', default='May', help=\"person name, e.g. Obama1, Obama2, May, Nadella, McStay\")\n",
        "    parser.add_argument('--driving_audio', default='./data/input/00083.wav', help=\"path to driving audio\")\n",
        "    parser.add_argument('--save_intermediates', default=0, help=\"whether to save intermediate results\")\n",
        "    \n",
        "    \n",
        "   \n",
        "    ############################### I/O Settings ##############################\n",
        "    # load config files\n",
        "    opt = parser.parse_args()\n",
        "    with open(join('./config/', opt.id + '.yaml')) as f:\n",
        "        config = yaml.load(f)\n",
        "    data_root = join('./data/', opt.id)\n",
        "    # create the results folder\n",
        "    audio_name = os.path.split(opt.driving_audio)[1][:-4]\n",
        "    save_root = join('./results/', opt.id, audio_name)\n",
        "    if not os.path.exists(save_root):\n",
        "        os.makedirs(save_root)\n",
        "       \n",
        "        \n",
        "\n",
        "    ############################ Hyper Parameters #############################\n",
        "    h, w, sr, FPS = 512, 512, 16000, 60\n",
        "    mouth_indices = np.concatenate([np.arange(4, 11), np.arange(46, 64)])\n",
        "    eye_brow_indices = [27, 65, 28, 68, 29, 67, 30, 66, 31, 72, 32, 69, 33, 70, 34, 71]\n",
        "    eye_brow_indices = np.array(eye_brow_indices, np.int32)    \n",
        "    \n",
        "\n",
        "    \n",
        "    ############################ Pre-defined Data #############################\n",
        "    mean_pts3d = np.load(join(data_root, 'mean_pts3d.npy'))  \n",
        "    fit_data = np.load(config['dataset_params']['fit_data_path'])\n",
        "    pts3d = np.load(config['dataset_params']['pts3d_path']) - mean_pts3d\n",
        "    trans = fit_data['trans'][:,:,0].astype(np.float32)\n",
        "    mean_translation = trans.mean(axis=0)\n",
        "    candidate_eye_brow = pts3d[10:, eye_brow_indices]\n",
        "    std_mean_pts3d = np.load(config['dataset_params']['pts3d_path']).mean(axis=0)\n",
        "    # candidates images    \n",
        "    img_candidates = []\n",
        "    for j in range(4):\n",
        "        output = imread(join(data_root, 'candidates', f'normalized_full_{j}.jpg'))\n",
        "        output = A.pytorch.transforms.ToTensor(normalize={'mean':(0.5,0.5,0.5), \n",
        "                                                          'std':(0.5,0.5,0.5)})(image=output)['image']\n",
        "        img_candidates.append(output)\n",
        "    img_candidates = torch.cat(img_candidates).unsqueeze(0).cuda()   \n",
        "    \n",
        "    # shoulders\n",
        "    shoulders = np.load(join(data_root, 'normalized_shoulder_points.npy'))\n",
        "    shoulder3D = np.load(join(data_root, 'shoulder_points3D.npy'))[1]\n",
        "    ref_trans = trans[1]    \n",
        "    \n",
        "    # camera matrix, we always use training set intrinsic parameters.\n",
        "    camera = utils.camera() \n",
        "    camera_intrinsic = np.load(join(data_root, 'camera_intrinsic.npy')).astype(np.float32)\n",
        "    APC_feat_database = np.load(join(data_root, 'APC_feature_base.npy'))\n",
        "    \n",
        "    # load reconstruction data\n",
        "    scale = sio.loadmat(join(data_root, 'id_scale.mat'))['scale'][0,0]\n",
        "    Audio2Mel_torch = audio_funcs.Audio2Mel(n_fft=512, hop_length=int(16000/120), win_length=int(16000/60), sampling_rate=16000, \n",
        "                                            n_mel_channels=80, mel_fmin=90, mel_fmax=7600.0).cuda()\n",
        "    \n",
        "    \n",
        "    \n",
        "    ########################### Experiment Settings ###########################\n",
        "    #### user config\n",
        "    use_LLE = config['model_params']['APC']['use_LLE']\n",
        "    Knear = config['model_params']['APC']['Knear']\n",
        "    LLE_percent = config['model_params']['APC']['LLE_percent']\n",
        "    headpose_sigma = config['model_params']['Headpose']['sigma']\n",
        "    Feat_smooth_sigma = config['model_params']['Audio2Mouth']['smooth']\n",
        "    Head_smooth_sigma = config['model_params']['Headpose']['smooth']\n",
        "    Feat_center_smooth_sigma, Head_center_smooth_sigma = 0, 0\n",
        "    AMP_method = config['model_params']['Audio2Mouth']['AMP'][0]\n",
        "    Feat_AMPs = config['model_params']['Audio2Mouth']['AMP'][1:]\n",
        "    rot_AMP, trans_AMP = config['model_params']['Headpose']['AMP']\n",
        "    shoulder_AMP = config['model_params']['Headpose']['shoulder_AMP']\n",
        "    save_feature_maps = config['model_params']['Image2Image']['save_input']\n",
        "    \n",
        "    #### common settings\n",
        "    Featopt = FeatureOptions().parse() \n",
        "    Headopt = HeadposeOptions().parse()\n",
        "    Renderopt = RenderOptions().parse()\n",
        "    Featopt.load_epoch = config['model_params']['Audio2Mouth']['ckp_path']\n",
        "    Headopt.load_epoch = config['model_params']['Headpose']['ckp_path']\n",
        "    Renderopt.dataroot = config['dataset_params']['root']\n",
        "    Renderopt.load_epoch = config['model_params']['Image2Image']['ckp_path']\n",
        "    Renderopt.size = config['model_params']['Image2Image']['size']\n",
        "    \n",
        "\n",
        "    \n",
        "    ############################# Load Models #################################\n",
        "    print('---------- Loading Model: APC-------------')\n",
        "    APC_model = APC_encoder(config['model_params']['APC']['mel_dim'],\n",
        "                            config['model_params']['APC']['hidden_size'],\n",
        "                            config['model_params']['APC']['num_layers'],\n",
        "                            config['model_params']['APC']['residual'])\n",
        "    APC_model.load_state_dict(torch.load(config['model_params']['APC']['ckp_path']), strict=False)\n",
        "    APC_model.cuda()\n",
        "    APC_model.eval()\n",
        "    print('---------- Loading Model: {} -------------'.format(Featopt.task))\n",
        "    Audio2Feature = create_model(Featopt)   \n",
        "    Audio2Feature.setup(Featopt)  \n",
        "    Audio2Feature.eval()     \n",
        "    print('---------- Loading Model: {} -------------'.format(Headopt.task))\n",
        "    Audio2Headpose = create_model(Headopt)    \n",
        "    Audio2Headpose.setup(Headopt)\n",
        "    Audio2Headpose.eval()              \n",
        "    if Headopt.feature_decoder == 'WaveNet':\n",
        "        Headopt.A2H_receptive_field = Audio2Headpose.Audio2Headpose.module.WaveNet.receptive_field\n",
        "    print('---------- Loading Model: {} -------------'.format(Renderopt.task))\n",
        "    facedataset = create_dataset(Renderopt) \n",
        "    Feature2Face = create_model(Renderopt)\n",
        "    Feature2Face.setup(Renderopt)   \n",
        "    Feature2Face.eval()\n",
        "    visualizer = Visualizer(Renderopt)\n",
        "    \n",
        "    \n",
        "    \n",
        "    ############################## Inference ##################################\n",
        "    print('Processing audio: {} ...'.format(audio_name)) \n",
        "    # read audio\n",
        "    audio, _ = librosa.load(opt.driving_audio, sr=sr)\n",
        "    total_frames = np.int32(audio.shape[0] / sr * FPS) \n",
        "    \n",
        "    \n",
        "    #### 1. compute APC features   \n",
        "    print('1. Computing APC features...')                    \n",
        "    mel80 = utils.compute_mel_one_sequence(audio)\n",
        "    mel_nframe = mel80.shape[0]\n",
        "    with torch.no_grad():\n",
        "        length = torch.Tensor([mel_nframe])\n",
        "        mel80_torch = torch.from_numpy(mel80.astype(np.float32)).cuda().unsqueeze(0)\n",
        "        hidden_reps = APC_model.forward(mel80_torch, length)[0]   # [mel_nframe, 512]\n",
        "        hidden_reps = hidden_reps.cpu().numpy()\n",
        "    audio_feats = hidden_reps\n",
        "            \n",
        "    \n",
        "    #### 2. manifold projection\n",
        "    if use_LLE:\n",
        "        print('2. Manifold projection...')\n",
        "        ind = utils.KNN_with_torch(audio_feats, APC_feat_database, K=Knear)\n",
        "        weights, feat_fuse = utils.compute_LLE_projection_all_frame(audio_feats, APC_feat_database, ind, audio_feats.shape[0])\n",
        "        audio_feats = audio_feats * (1-LLE_percent) + feat_fuse * LLE_percent\n",
        "    \n",
        "    \n",
        "    #### 3. Audio2Mouth\n",
        "    print('3. Audio2Mouth inference...')\n",
        "    pred_Feat = Audio2Feature.generate_sequences(audio_feats, sr, FPS, fill_zero=True, opt=Featopt)\n",
        "    \n",
        "    \n",
        "    #### 4. Audio2Headpose\n",
        "    print('4. Headpose inference...')\n",
        "    # set history headposes as zero\n",
        "    pre_headpose = np.zeros(Headopt.A2H_wavenet_input_channels, np.float32)\n",
        "    pred_Head = Audio2Headpose.generate_sequences(audio_feats, pre_headpose, fill_zero=True, sigma_scale=0.3, opt=Headopt)\n",
        "       \n",
        "    \n",
        "    #### 5. Post-Processing \n",
        "    print('5. Post-processing...')\n",
        "    nframe = min(pred_Feat.shape[0], pred_Head.shape[0])\n",
        "    pred_pts3d = np.zeros([nframe, 73, 3])\n",
        "    pred_pts3d[:, mouth_indices] = pred_Feat.reshape(-1, 25, 3)[:nframe]\n",
        "    \n",
        "    ## mouth\n",
        "    pred_pts3d = utils.landmark_smooth_3d(pred_pts3d, Feat_smooth_sigma, area='only_mouth')\n",
        "    pred_pts3d = utils.mouth_pts_AMP(pred_pts3d, True, AMP_method, Feat_AMPs)\n",
        "    pred_pts3d = pred_pts3d + mean_pts3d\n",
        "    pred_pts3d = utils.solve_intersect_mouth(pred_pts3d)  # solve intersect lips if exist\n",
        "\n",
        "    ## headpose\n",
        "    pred_Head[:, 0:3] *= rot_AMP\n",
        "    pred_Head[:, 3:6] *= trans_AMP\n",
        "    pred_headpose = utils.headpose_smooth(pred_Head[:,:6], Head_smooth_sigma).astype(np.float32)\n",
        "    pred_headpose[:, 3:] += mean_translation\n",
        "    pred_headpose[:, 0] += 180\n",
        "    \n",
        "    ## compute projected landmarks\n",
        "    pred_landmarks = np.zeros([nframe, 73, 2], dtype=np.float32)\n",
        "    final_pts3d = np.zeros([nframe, 73, 3], dtype=np.float32)\n",
        "    final_pts3d[:] = std_mean_pts3d.copy()\n",
        "    final_pts3d[:, 46:64] = pred_pts3d[:nframe, 46:64]\n",
        "    for k in tqdm(range(nframe)):\n",
        "        ind = k % candidate_eye_brow.shape[0]\n",
        "        final_pts3d[k, eye_brow_indices] = candidate_eye_brow[ind] + mean_pts3d[eye_brow_indices]\n",
        "        pred_landmarks[k], _, _ = utils.project_landmarks(camera_intrinsic, camera.relative_rotation, \n",
        "                                                          camera.relative_translation, scale, \n",
        "                                                          pred_headpose[k], final_pts3d[k]) \n",
        "    \n",
        "    ## Upper Body Motion\n",
        "    pred_shoulders = np.zeros([nframe, 18, 2], dtype=np.float32)\n",
        "    pred_shoulders3D = np.zeros([nframe, 18, 3], dtype=np.float32)\n",
        "    for k in range(nframe):\n",
        "        diff_trans = pred_headpose[k][3:] - ref_trans\n",
        "        pred_shoulders3D[k] = shoulder3D + diff_trans * shoulder_AMP\n",
        "        # project\n",
        "        project = camera_intrinsic.dot(pred_shoulders3D[k].T)\n",
        "        project[:2, :] /= project[2, :]  # divide z\n",
        "        pred_shoulders[k] = project[:2, :].T\n",
        "    \n",
        "\n",
        "    #### 6. Image2Image translation & Save resuls\n",
        "    print('6. Image2Image translation & Saving results...')\n",
        "    for ind in tqdm(range(0, nframe), desc='Image2Image translation inference'):\n",
        "        # feature_map: [input_nc, h, w]\n",
        "        current_pred_feature_map = facedataset.dataset.get_data_test_mode(pred_landmarks[ind], \n",
        "                                                                          pred_shoulders[ind], \n",
        "                                                                          facedataset.dataset.image_pad)\n",
        "        input_feature_maps = current_pred_feature_map.unsqueeze(0).cuda()\n",
        "        pred_fake = Feature2Face.inference(input_feature_maps, img_candidates) \n",
        "        # save results\n",
        "        visual_list = [('pred', util.tensor2im(pred_fake[0]))]\n",
        "        if save_feature_maps:\n",
        "            visual_list += [('input', np.uint8(current_pred_feature_map[0].cpu().numpy() * 255))]\n",
        "        visuals = OrderedDict(visual_list)\n",
        "        visualizer.save_images(save_root, visuals, str(ind+1))\n",
        "    \n",
        "\n",
        "    ## make videos\n",
        "    # generate corresponding audio, reused for all results\n",
        "    tmp_audio_path = join(save_root, 'tmp.wav')\n",
        "    tmp_audio_clip = audio[ : np.int32(nframe * sr / FPS)]\n",
        "    #librosa.output.write_wav(tmp_audio_path, tmp_audio_clip, sr)\n",
        "    sf.write(tmp_audio_path, tmp_audio_clip, sr)\n",
        "    \n",
        "    \n",
        "    final_path = join(save_root, audio_name + '.avi')\n",
        "    write_video_with_audio(tmp_audio_path, final_path, 'pred_')\n",
        "    feature_maps_path = join(save_root, audio_name + '_feature_maps.avi')\n",
        "    write_video_with_audio(tmp_audio_path, feature_maps_path, 'input_')\n",
        "    \n",
        "    if os.path.exists(tmp_audio_path):\n",
        "        os.remove(tmp_audio_path)\n",
        "    if not opt.save_intermediates:\n",
        "        _img_paths = list(map(lambda x:str(x), list(Path(save_root).glob('*.jpg'))))\n",
        "        for i in tqdm(range(len(_img_paths)), desc='deleting intermediate images'):\n",
        "            os.remove(_img_paths[i])\n",
        "    \n",
        "    \n",
        "    print('Finish!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "sag5dYkLgydS"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Upload audio</font></b>\n",
        "%cd /content/LiveSpeechPortraits\n",
        "uploaded = files.upload()\n",
        "for fn in uploaded.keys():\n",
        "  print('User uploaded file \"{name}\" with length {length} bytes'.format(\n",
        "      name=fn, length=len(uploaded[fn])))\n",
        "  \n",
        "audio_type = fn.split('.')[-1]\n",
        "os.rename(fn, 'audio.'+audio_type)\n",
        "fn = '/content/LiveSpeechPortraits/audio.'+audio_type\n",
        "clear_output()"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bUk1E8S8RBYJ",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Process</font></b>\n",
        "model = \"May\" #@param [\"May\", \"McStay\", \"Obama1\", \"Obama2\", \"Nadella\"]\n",
        "!python demo.py --id $model --driving_audio $fn\n",
        "\n",
        "video = f'/content/LiveSpeechPortraits/results/{model}/audio/audio.avi'\n",
        "f_map = f'/content/LiveSpeechPortraits/results/{model}/audio/audio_feature_maps.avi'\n",
        "new_vid = f'/content/LiveSpeechPortraits/results/{model}/audio/side.mp4'\n",
        "\n",
        "clear_output()"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "XuGWK0q8i4kq"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Download video</font></b>\n",
        "\n",
        "files.download(f'/content/LiveSpeechPortraits/results/{model}/audio/audio.avi')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S5MDqTu9i--p",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Download feature maps</font></b>\n",
        "\n",
        "files.download(f'/content/LiveSpeechPortraits/results/{model}/audio/audio_feature_maps.avi')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_yj2eXIIjtnk",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Download side-by-syde</font></b>\n",
        "!ffmpeg -y -i $f_map -i $video -filter_complex \"[0:v:0]pad=iw*2:ih[bg]; [bg][1:v:0]overlay=w\" $new_vid\n",
        "\n",
        "files.download(new_vid)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}