GPEN_DFL_Colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GPEN DFL Colab",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "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/GPEN_DFL_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZRnUFdLSJ_WD"
      },
      "source": [
        "<b><font size=\"+4\">GPEN for DeepFaceLab</font></b>\n",
        "\n",
        "<b><font size=\"+2\">Based on:</font></b>\n",
        "\n",
        "**GitHub repository**: [GPEN](https://github.com/yangxy/GPEN)\n",
        "\n",
        "Article: [GAN Prior Embedded Network for Blind Face Restoration in the Wild](https://arxiv.org/pdf/2105.06070.pdf)\n",
        "\n",
        "Creator: **[yangxy](https://github.com/yangxy)**\n",
        "\n",
        "This Сolab was created by [Bomze](https://github.com/tg-bomze/)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O2nm3nGkLwJj",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font size=\"+3\">0) Check GPU</font></b>\n",
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GETEjyBz4cFj",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font size=\"+3\">1) Prepare Colab machine</font></b>\n",
        "!pip install ffmpeg-python\n",
        "#!pip install torch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2\n",
        "#!pip install pip install opencv-python\n",
        "\n",
        "from IPython.display import clear_output\n",
        "from google.colab import files\n",
        "!git clone https://github.com/yangxy/GPEN\n",
        "\n",
        "!wget https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth && mv RetinaFace-R50.pth GPEN/weights/\n",
        "!wget https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-512.pth && mv GPEN-512.pth GPEN/weights/\n",
        "#!wget https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-1024-Color.pth && mv GPEN-1024-Color.pth GPEN/weights/\n",
        "\n",
        "import os\n",
        "import cv2\n",
        "import ffmpeg\n",
        "import sys\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "%matplotlib inline\n",
        "%cd /content/GPEN\n",
        "\n",
        "!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip\n",
        "!sudo unzip ninja-linux.zip -d /usr/local/bin/\n",
        "!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force\n",
        "clear_output()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0RF8vZUDL7Ma",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font size=\"+3\">2) Simple settings</font></b>\n",
        "%%writefile /content/GPEN/main_file.py\n",
        "import os\n",
        "import cv2\n",
        "import glob\n",
        "import time\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "from tqdm import tqdm\n",
        "import __init_paths\n",
        "from retinaface.retinaface_detection import RetinaFaceDetection\n",
        "from face_model.face_gan import FaceGAN\n",
        "from align_faces import warp_and_crop_face, get_reference_facial_points\n",
        "from skimage import transform as tf\n",
        "\n",
        "'''class FaceColorization(object):\n",
        "    def __init__(self, base_dir='./', size=1024, model=None, channel_multiplier=2):\n",
        "        self.facegan = FaceGAN(base_dir, size, model, channel_multiplier)\n",
        "    # make sure the face image is well aligned. Please refer to face_enhancement.py\n",
        "    def process(self, gray):\n",
        "        return self.facegan.process(gray)#'''\n",
        "\n",
        "class FaceEnhancement(object):\n",
        "    def __init__(self, base_dir='./', size=512, model=None, channel_multiplier=2):\n",
        "        self.facedetector = RetinaFaceDetection(base_dir)\n",
        "        self.facegan = FaceGAN(base_dir, size, model, channel_multiplier)\n",
        "        self.size = size\n",
        "        self.threshold = 0.9\n",
        "\n",
        "        # the mask for pasting restored faces back\n",
        "        self.mask = np.zeros((512, 512), np.float32)\n",
        "        cv2.rectangle(self.mask, (26, 26), (486, 486), (1, 1, 1), -1, cv2.LINE_AA)\n",
        "        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)\n",
        "        self.mask = cv2.GaussianBlur(self.mask, (101, 101), 11)\n",
        "\n",
        "        self.kernel = np.array((\n",
        "                [0.0625, 0.125, 0.0625],\n",
        "                [0.125, 0.25, 0.125],\n",
        "                [0.0625, 0.125, 0.0625]), dtype=\"float32\")\n",
        "\n",
        "        # get the reference 5 landmarks position in the crop settings\n",
        "        default_square = True\n",
        "        inner_padding_factor = 0.25\n",
        "        outer_padding = (0, 0)\n",
        "        self.reference_5pts = get_reference_facial_points(\n",
        "                (self.size, self.size), inner_padding_factor, outer_padding, default_square)\n",
        "\n",
        "    def process(self, img):\n",
        "        facebs, landms = self.facedetector.detect(img)\n",
        "        \n",
        "        orig_faces, enhanced_faces = [], []\n",
        "        height, width = img.shape[:2]\n",
        "        full_mask = np.zeros((height, width), dtype=np.float32)\n",
        "        full_img = np.zeros(img.shape, dtype=np.uint8)\n",
        "\n",
        "        for i, (faceb, facial5points) in enumerate(zip(facebs, landms)):\n",
        "            if faceb[4]<self.threshold: continue\n",
        "            fh, fw = (faceb[3]-faceb[1]), (faceb[2]-faceb[0])\n",
        "\n",
        "            facial5points = np.reshape(facial5points, (2, 5))\n",
        "\n",
        "            of, tfm_inv = warp_and_crop_face(img, facial5points, reference_pts=self.reference_5pts, crop_size=(self.size, self.size))\n",
        "            \n",
        "            # enhance the face\n",
        "            ef = self.facegan.process(of)\n",
        "            \n",
        "            orig_faces.append(of)\n",
        "            enhanced_faces.append(ef)\n",
        "            \n",
        "            tmp_mask = self.mask\n",
        "            tmp_mask = cv2.resize(tmp_mask, ef.shape[:2])\n",
        "            tmp_mask = cv2.warpAffine(tmp_mask, tfm_inv, (width, height), flags=3)\n",
        "\n",
        "            if min(fh, fw)<100: # gaussian filter for small faces\n",
        "                ef = cv2.filter2D(ef, -1, self.kernel)\n",
        "            \n",
        "            tmp_img = cv2.warpAffine(ef, tfm_inv, (width, height), flags=3)\n",
        "\n",
        "            mask = tmp_mask - full_mask\n",
        "            full_mask[np.where(mask>0)] = tmp_mask[np.where(mask>0)]\n",
        "            full_img[np.where(mask>0)] = tmp_img[np.where(mask>0)]\n",
        "\n",
        "        full_mask = full_mask[:, :, np.newaxis]\n",
        "        img = cv2.convertScaleAbs(img*(1-full_mask) + full_img*full_mask)\n",
        "\n",
        "        return img, orig_faces, enhanced_faces\n",
        "        \n",
        "\n",
        "if __name__=='__main__':\n",
        "    indir = '/content/frames'\n",
        "    outdir = '/content/result'\n",
        "    scale = '1' #@param {type:\"string\"}\n",
        "    #@markdown scale - how many times to enlarge the input image.\n",
        "    os.makedirs(outdir, exist_ok=True)\n",
        "\n",
        "    faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2)\n",
        "    #facecolorizer = FaceColorization(size=1024, model='GPEN-1024-Color', channel_multiplier=2)\n",
        "\n",
        "    files = sorted(glob.glob(os.path.join(indir, '*.*g')))\n",
        "    for n, file in tqdm(enumerate(files[:])):\n",
        "        filename = os.path.basename(file)\n",
        "        im = cv2.imread(file, cv2.IMREAD_COLOR)\n",
        "        \n",
        "        # Enchance\n",
        "        if not isinstance(im, np.ndarray): print(filename, 'error'); continue\n",
        "        im = cv2.resize(im, (0,0), fx=int(scale), fy=int(scale))\n",
        "        im, orig_faces, enhanced_faces = faceenhancer.process(im)\n",
        "        cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'.png'), im)#'''\n",
        "\n",
        "        # Colorize\n",
        "        '''colorf = facecolorizer.process(im)\n",
        "        colorf = cv2.resize(colorf, (im.shape[1], im.shape[0]))\n",
        "        cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'.png'), colorf)#'''"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pbFO8GCI-DCr",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font size=\"+3\">3) Mount Google Drive</font></b>\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OZ7-E-vRSYar",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font size=\"+3\">4) Enhance!</font></b>\n",
        "\n",
        "import os\n",
        "\n",
        "working_dir = \"/content/drive/MyDrive/\" #@param {type:\"string\"}\n",
        "aligned_zip = True #@param {type:\"boolean\"}\n",
        "name_of_zip_or_folder = \"aligned.zip\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Pick a folder from your Google Drive.\n",
        "\n",
        "#@markdown **working_dir** should contain **aligned** subfolder or **aligned.zip** file with JPG (or PNG) images that were extracted by DeepFaceLab *'faceset extract'* scripts.\n",
        "\n",
        "#@markdown Results will be saved in **result** subfolder or **result.zip** file in **working_dir**.\n",
        "\n",
        "#@markdown If working not in zip mode existing files will be skipped so you can easily continue in case of disconnect.\n",
        "\n",
        "path_to_zip_or_folder = f'/content/{name_of_zip_or_folder}'\n",
        "!rm -r '/content/frames' '/content/result' '/content/result.zip'\n",
        "%cd /content/\n",
        "if aligned_zip:\n",
        "  zip_path = os.path.join(working_dir,name_of_zip_or_folder)\n",
        "  if not os.path.isfile(zip_path):\n",
        "    print('\"{}\" not found'.format(zip_path))\n",
        "    sys.exit()\n",
        "  !rsync -ah --progress \"$zip_path\" '/content/'\n",
        "  !mkdir -p '/content/frames'\n",
        "  !unzip $path_to_zip_or_folder -d /content/frames\n",
        "  if len(os.listdir('/content/frames')) == 1:\n",
        "    !rm -rf /content/frames/*\n",
        "    !mkdir /content/temp\n",
        "    !unzip $path_to_zip_or_folder -d /content/temp\n",
        "    all_files = \"/content/temp/\" + os.listdir(\"/content/temp\")[0] + \"/*\"\n",
        "    !mv $all_files /content/frames\n",
        "  !rm -rf $path_to_zip_or_folder /content/frames/__MACOSX /content/temp\n",
        "else:\n",
        "  in_path = os.path.join(working_dir,name_of_zip_or_folder)\n",
        "  out_path = os.path.join(working_dir,'result')\n",
        "  if not os.path.isdir(in_path):\n",
        "    print('\"{}\" not found'.format(in_path))\n",
        "    sys.exit()\n",
        "  !rsync -ah --progress \"$in_path\" '/content'\n",
        "  !mv $path_to_zip_or_folder /content/frames\n",
        "\n",
        "!mkdir /content/result\n",
        "%cd /content/GPEN\n",
        "clear_output()\n",
        "print('Start Enhancment')\n",
        "!python main_file.py\n",
        "\n",
        "%cd /content/\n",
        "if aligned_zip:\n",
        "  !rm /content/result.zip\n",
        "  !7z a /content/result.zip /content/result/*\n",
        "  output_zip = os.path.join(working_dir,'result.zip')\n",
        "  !rsync -ah --progress /content/result.zip \"$output_zip\"\n",
        "  clear_output()\n",
        "  print('Results saved to \"{}\"'.format(output_zip))\n",
        "else:\n",
        "  !rsync -ah --progress '/content/result' \"$working_dir\"\n",
        "  clear_output()\n",
        "  print('Results saved to \"{}\"'.format(out_path))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}