AnimeFaceAlign.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "AnimeFaceAlign.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyMw1J0TpEqdUsa6vRI9P6PO",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/RaghavaDhanya/AnimeFaceAlign/blob/main/AnimeFaceAlign.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UgOX4XLViBEv"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/RaghavaDhanya/AnimeFaceAlign/blob/main/AnimeFaceAlign.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JHIgwMTdXsP8",
        "outputId": "73543996-fde1-4e47-80ad-6b70e6874690"
      },
      "source": [
        "!git clone https://github.com/kanosawa/anime_face_landmark_detection.git"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'anime_face_landmark_detection'...\n",
            "remote: Enumerating objects: 35, done.\u001b[K\n",
            "remote: Counting objects: 100% (35/35), done.\u001b[K\n",
            "remote: Compressing objects: 100% (24/24), done.\u001b[K\n",
            "remote: Total 35 (delta 13), reused 26 (delta 7), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (35/35), done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2hcUB3V2YQD1",
        "outputId": "c74c4b8c-9502-44cf-f585-bd86d6a3f9aa"
      },
      "source": [
        "%cd anime_face_landmark_detection/"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/anime_face_landmark_detection\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c2n_4dJXYRup",
        "outputId": "85e5d1e5-f420-43b9-8a47-e3aa3e60a475"
      },
      "source": [
        "!gdown https://drive.google.com/uc?id=1NckKw7elDjQTllRxttO87WY7cnQwdMqz"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1NckKw7elDjQTllRxttO87WY7cnQwdMqz\n",
            "To: /content/anime_face_landmark_detection/checkpoint_landmark_191116.pth.tar\n",
            "16.7MB [00:00, 63.3MB/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MZWWp6JaYdHV",
        "outputId": "b90bc9e4-2575-49a0-bc7a-e407e6c33d2f"
      },
      "source": [
        "!wget https://raw.githubusercontent.com/nagadomi/lbpcascade_animeface/master/lbpcascade_animeface.xml"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2021-09-18 18:31:48--  https://raw.githubusercontent.com/nagadomi/lbpcascade_animeface/master/lbpcascade_animeface.xml\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.108.133, 185.199.109.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 246945 (241K) [text/plain]\n",
            "Saving to: ‘lbpcascade_animeface.xml’\n",
            "\n",
            "lbpcascade_animefac 100%[===================>] 241.16K  --.-KB/s    in 0.03s   \n",
            "\n",
            "2021-09-18 18:31:48 (7.72 MB/s) - ‘lbpcascade_animeface.xml’ saved [246945/246945]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hbQ3FUZTYpyg"
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "from torchvision import transforms\n",
        "import cv2\n",
        "from PIL import Image, ImageDraw, ImageFont\n",
        "\n",
        "from CFA import CFA\n",
        "# import animeface\n",
        "\n",
        "\n",
        "# param\n",
        "\n",
        "def get_landmarks(img, face_detector, landmark_detector, img_width=128, num_landmark = 24):\n",
        "# transform\n",
        "    normalize   = transforms.Normalize(mean=[0.5, 0.5, 0.5],\n",
        "                                    std=[0.5, 0.5, 0.5])\n",
        "    train_transform = [transforms.ToTensor(), normalize]\n",
        "    train_transform = transforms.Compose(train_transform)\n",
        "\n",
        "    # input image & detect face\n",
        "    # img = cv2.imread(input_img_name)\n",
        "\n",
        "    faces = face_detector.detectMultiScale(cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB),scaleFactor =1.05)\n",
        "    # draw = ImageDraw.Draw(img)\n",
        "\n",
        "    for x_, y_, w_, h_ in faces:\n",
        "        # adjust face size\n",
        "        x = max(x_ - w_ / 8, 0)\n",
        "        rx = min(x_ + w_ * 9 / 8, img.width)\n",
        "        y = max(y_ - h_ / 4, 0)\n",
        "        by = y_ + h_\n",
        "        w = rx - x\n",
        "        h = by - y\n",
        "\n",
        "        # draw result of face detection\n",
        "        # draw.rectangle((x, y, x + w, y + h), outline=(0, 0, 255), width=3)\n",
        "\n",
        "        # transform image\n",
        "        img_tmp = img.crop((x, y, x+w, y+h))\n",
        "        img_tmp = img_tmp.resize((img_width, img_width), Image.BICUBIC)\n",
        "        img_tmp = train_transform(img_tmp)\n",
        "        img_tmp = img_tmp.unsqueeze(0).cuda()\n",
        "\n",
        "        # estimate heatmap\n",
        "        heatmaps = landmark_detector(img_tmp)\n",
        "        heatmaps = heatmaps[-1].cpu().detach().numpy()[0]\n",
        "        landmarks = []\n",
        "        # calculate landmark position\n",
        "        for i in range(num_landmark):\n",
        "            heatmaps_tmp = cv2.resize(heatmaps[i], (img_width, img_width), interpolation=cv2.INTER_CUBIC)\n",
        "            landmark = np.unravel_index(np.argmax(heatmaps_tmp), heatmaps_tmp.shape)\n",
        "            landmark_y = landmark[0] * h / img_width\n",
        "            landmark_x = landmark[1] * w / img_width\n",
        "            landmarks.append([x+landmark_x, y+landmark_y])\n",
        "        yield landmarks\n",
        "\n",
        "# output image\n",
        "# img.save('output.png')"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nuYfFW3rgYqO"
      },
      "source": [
        "import scipy.ndimage\n",
        "from pathlib import Path\n",
        "\n",
        "def align_image(img_path, output_size=1024, transform_size=4096, enable_padding=True, x_scale=1, y_scale=1, em_scale=0.1, alpha=False):\n",
        "    checkpoint_name = 'checkpoint_landmark_191116.pth.tar'\n",
        "    img_orig =Image.open(img_path) \n",
        "    num_landmark = 24\n",
        "\n",
        "    # detector\n",
        "    face_detector = cv2.CascadeClassifier('lbpcascade_animeface.xml')\n",
        "    landmark_detector = CFA(output_channel_num=num_landmark + 1, checkpoint_name=checkpoint_name).cuda()\n",
        "    for i, landmarks in enumerate(get_landmarks(img_orig, face_detector, landmark_detector, num_landmark=num_landmark)):\n",
        "        img = img_orig.copy()\n",
        "\n",
        "        lm = np.array(landmarks)\n",
        "        face_outline = lm[[0, 1, 2]]\n",
        "        left_eye_brow = lm[[3, 4, 5]]\n",
        "        right_eye_brow = lm[[6, 7, 8]]\n",
        "        nose = lm[[9]]\n",
        "        lm_eye_left = lm[[10, 11, 12, 13, 14]]\n",
        "        lm_eye_right = lm[[15, 16, 17, 18, 19]]\n",
        "        mouth = lm[[20, 21, 22, 23]]\n",
        "\n",
        "        mouth_left_point = lm[20]\n",
        "        mount_right_point = lm[22]\n",
        "\n",
        "\n",
        "        eye_left     = np.mean(lm_eye_left, axis=0)\n",
        "        eye_right    = np.mean(lm_eye_right, axis=0)\n",
        "        eye_avg      = (eye_left + eye_right) * 0.5\n",
        "        eye_to_eye   = eye_right - eye_left\n",
        "        mouth_avg    = (mouth_left_point + mount_right_point) * 0.5\n",
        "        eye_to_mouth = mouth_avg - eye_avg\n",
        "\n",
        "        # Choose oriented crop rectangle.\n",
        "        x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]\n",
        "        x /= np.hypot(*x)\n",
        "        x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)\n",
        "        x *= x_scale\n",
        "        y = np.flipud(x) * [-y_scale, y_scale]\n",
        "        c = eye_avg + eye_to_mouth * em_scale\n",
        "        quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])\n",
        "        qsize = np.hypot(*x) * 2\n",
        "\n",
        "        # Shrink.\n",
        "        shrink = int(np.floor(qsize / output_size * 0.5))\n",
        "        if shrink > 1:\n",
        "            rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink)))\n",
        "            img = img.resize(rsize, Image.ANTIALIAS)\n",
        "            quad /= shrink\n",
        "            qsize /= shrink\n",
        "\n",
        "        # Crop.\n",
        "        border = max(int(np.rint(qsize * 0.1)), 3)\n",
        "        crop = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1]))))\n",
        "        crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]), min(crop[3] + border, img.size[1]))\n",
        "        if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:\n",
        "            img = img.crop(crop)\n",
        "            quad -= crop[0:2]\n",
        "\n",
        "        # Pad.\n",
        "        pad = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1]))))\n",
        "        pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0), max(pad[3] - img.size[1] + border, 0))\n",
        "        if enable_padding and max(pad) > border - 4:\n",
        "            pad = np.maximum(pad, int(np.rint(qsize * 0.3)))\n",
        "            img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')\n",
        "            h, w, _ = img.shape\n",
        "            y, x, _ = np.ogrid[:h, :w, :1]\n",
        "            mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w-1-x) / pad[2]), 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h-1-y) / pad[3]))\n",
        "            blur = qsize * 0.02\n",
        "            img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)\n",
        "            img += (np.median(img, axis=(0,1)) - img) * np.clip(mask, 0.0, 1.0)\n",
        "            img = np.uint8(np.clip(np.rint(img), 0, 255))\n",
        "            if alpha:\n",
        "                mask = 1-np.clip(3.0 * mask, 0.0, 1.0)\n",
        "                mask = np.uint8(np.clip(np.rint(mask*255), 0, 255))\n",
        "                img = np.concatenate((img, mask), axis=2)\n",
        "                img = Image.fromarray(img, 'RGBA')\n",
        "            else:\n",
        "                img = Image.fromarray(img, 'RGB')\n",
        "            quad += pad[:2]\n",
        "\n",
        "        # Transform.\n",
        "        img = img.transform((transform_size, transform_size), Image.QUAD, (quad + 0.5).flatten(), Image.BILINEAR)\n",
        "        if output_size < transform_size:\n",
        "            img = img.resize((output_size, output_size), Image.ANTIALIAS)\n",
        "\n",
        "        # Save aligned image.\n",
        "        img.save(f'{Path(img_path).stem}_{i}.png', 'PNG')"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "riK8BhF4_3u8"
      },
      "source": [
        "align_image('/content/anime_face_landmark_detection/test.png')"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U38fL8SCOEMC"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}