{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "from anime_face_detector import create_detector\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_deg(arr):\n",
    "    rad = math.atan2(arr[3]-arr[1],arr[2]-arr[0])\n",
    "    PI = math.pi\n",
    "    deg = (rad*180)/PI\n",
    "    return deg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detector = create_detector('yolov3', device='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gg = Image.open('gg.png')\n",
    "# ggdraw = ImageDraw.Draw(gg)\n",
    "# ggdraw.rectangle((5,5,gg.width-5,gg.height-5), outline=(255, 0,0), width=5)\n",
    "gg = cv2.imread('gg.png', cv2.IMREAD_UNCHANGED)\n",
    "gg = cv2.cvtColor(gg, cv2.COLOR_BGRA2RGBA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(604, 1074, 3)\n"
     ]
    }
   ],
   "source": [
    "img = cv2.imread('test.png')\n",
    "preds = detector(img)\n",
    "img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n",
    "# draw = ImageDraw.Draw(image)\n",
    "\n",
    "# for face in preds:\n",
    "#     draw.rectangle((face['bbox'][0], face['bbox'][1], face['bbox'][2], face['bbox'][3]), outline=(255, 0, 0), width=5)\n",
    "#     x = face['bbox'][0]\n",
    "#     y = face['bbox'][1]\n",
    "#     for i, point in enumerate(face['keypoints']):\n",
    "#         # draw.ellipse((point[0]-2, point[1]-2, point[0]+2, point[1]+2), fill=(255, 0, 0))\n",
    "#         draw.text((point[0], point[1]), str(i), font=ImageFont.truetype('arial.ttf', 10), fill=(255, 0, 0))\n",
    "\n",
    "for face in preds:\n",
    "    points = face['keypoints']\n",
    "    color = img[int(points[27][1]), int(points[27][0])+10]\n",
    "    polygon = np.array([\n",
    "        [points[0][0], points[0][1]],\n",
    "        [points[1][0], points[1][1]],\n",
    "        [points[2][0], points[2][1]],\n",
    "        [points[3][0], points[4][1]],\n",
    "        [points[4][0], points[4][1]],\n",
    "        [points[10][0], points[10][1]],\n",
    "        [points[9][0], points[9][1]],\n",
    "        [points[8][0], points[8][1]],\n",
    "        [points[7][0], points[7][1]],\n",
    "        [points[6][0], points[6][1]],\n",
    "        [points[5][0], points[5][1]]\n",
    "    ], np.int32)\n",
    "    cv2.fillConvexPoly(img, polygon, color=(int(color[0]), int(color[1]), int(color[2]), 255))\n",
    "    deg = get_deg([points[0][0], points[0][1], points[4][0], points[4][1]])\n",
    "    rotated = gg.copy()\n",
    "    resize = math.sqrt((points[10][0] - points[5][0])**2 + (points[10][1] - points[5][1])**2)\n",
    "    rotated = cv2.resize(rotated, (int(resize), int(resize*1.12)))\n",
    "    matrix = cv2.getPerspectiveTransform(\n",
    "        np.float32([[0, 0], [rotated.shape[0],0], [0, rotated.shape[1]], [rotated.shape[0],rotated.shape[1]]]),\n",
    "        np.float32([[points[5][0], points[5][1]], [points[10][0], points[10][1]], [points[1][0], points[1][1]], [points[3][0], points[3][1]]]))\n",
    "    rotated = cv2.warpPerspective(rotated, matrix, (img.shape[1], img.shape[0]))\n",
    "\n",
    "    alpha = rotated[:, :, 3] / 255.\n",
    "    for i in range(3):\n",
    "        w, h = img.shape[:2]\n",
    "        rw, rh = rotated.shape[:2]\n",
    "        pointx, pointy = points[5][:2]\n",
    "        pointx, pointy = int(pointx), int(pointy)\n",
    "        img[:, :, i] = (1. - alpha) * img[0:, 0:, i] + alpha * rotated[:, :, i]\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.asarray(img))\n",
    "cv2.imwrite('result.png', cv2.cvtColor(img, cv2.COLOR_RGBA2BGR))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('bot')",
   "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.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "1230fa9187aff02e6ebfc79b73c9c8422b1bc4886baebd37c443f3278ff8d769"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}