1.stylegan_preprocess_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdf4e2b9-4831-4695-809c-e4b74cfe03a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# based on https://github.com/happy-jihye/FFHQ-Alignment/blob/master/Anime-Face-Alignment\n",
    "# https://github.com/hysts/anime-face-detector/blob/main/demo.ipynb\n",
    "\n",
    "!sudo apt install ffmpeg\n",
    "!pip install face-alignment\n",
    "!pip install opencv-python\n",
    "!git clone https://github.com/NVlabs/stylegan3.git\n",
    "\n",
    "# anime-face-detector \n",
    "!pip install openmim\n",
    "!mim install mmcv-full mmdet mmpose -y\n",
    "!pip install anime-face-detector --no-dependencies\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdf2c07d-98ee-46cb-9a4c-cbc199498384",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from glob import glob\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import scipy.ndimage\n",
    "import cv2\n",
    "import PIL\n",
    "from PIL import Image\n",
    "from shutil import copyfile\n",
    "from IPython.display import display\n",
    "import face_alignment\n",
    "import anime_face_detector\n",
    "\n",
    "videoDir = \"Arcane\"\n",
    "frameDir = \"frames\"\n",
    "alignedDir = \"alignedFace\"\n",
    "filteredDir = \"filteredFace\"\n",
    "preprocessedDir = \"preprocessedFace\"\n",
    "dataZip= \"arcaneFilteredData.zip\"\n",
    "\n",
    "for i in [videoDir,frameDir,alignedDir,filteredDir,preprocesseddDir]:\n",
    "    os.makedirs(i, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdb50f7f-345a-4d93-a0bb-5eabf30885db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # get frames from video\n",
    "# videoList=glob(videoDir+\"/*.mp4\")\n",
    "\n",
    "# # get 2 frame per sec, best jpg quality \n",
    "# for file in videoList:\n",
    "#     name=Path(file).stem\n",
    "#     !ffmpeg -i \"$file\" -r 2 -q:v 1 -qmin 1 -qmax 1 \"$frameDir\"/\"$name\"_%04d.jpg\n",
    "# PIL.Image.open(glob(frameDir+\"/*.jpg\")[4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d3344b8-c23d-447c-9a27-34d574fe6b74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_align(src_file, dst_file, face_landmarks, output_size=256, transform_size=1024, enable_padding=True, use_landmark_28=False):\n",
    "    # Align function from FFHQ dataset pre-processing step\n",
    "    # https://github.com/NVlabs/ffhq-dataset/blob/master/download_ffhq.py\n",
    "\n",
    "\n",
    "    if(use_landmark_28==False):\n",
    "        lm = np.array(face_landmarks)\n",
    "        lm_chin          = lm[0  : 17, :2]  # left-right\n",
    "        lm_eyebrow_left  = lm[17 : 22, :2]  # left-right\n",
    "        lm_eyebrow_right = lm[22 : 27, :2]  # left-right\n",
    "        lm_nose          = lm[27 : 31, :2]  # top-down\n",
    "        lm_nostrils      = lm[31 : 36, :2]  # top-down\n",
    "        lm_eye_left      = lm[36 : 42, :2]  # left-clockwise\n",
    "        lm_eye_right     = lm[42 : 48, :2]  # left-clockwise\n",
    "        lm_mouth_outer   = lm[48 : 60, :2]  # left-clockwise\n",
    "        lm_mouth_inner   = lm[60 : 68, :2]  # left-clockwise\n",
    "        mouth_left   = lm_mouth_outer[0]\n",
    "        mouth_right  = lm_mouth_outer[6]\n",
    "        \n",
    "    else:\n",
    "        lm = np.array(face_landmarks)\n",
    "        lm_eye_left      = lm[11 : 17, :2]  # left-clockwise\n",
    "        lm_eye_right     = lm[17 : 23, :2]  # left-clockwise\n",
    "        mouth_left   = lm[24, :2]\n",
    "        mouth_right  = lm[26, :2]\n",
    "\n",
    "\n",
    "    # Calculate auxiliary vectors.\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 + mouth_right) * 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",
    "    y = np.flipud(x) * [-1, 1]\n",
    "    c = eye_avg + eye_to_mouth * 0.1\n",
    "    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])\n",
    "    qsize = np.hypot(*x) * 2\n",
    "\n",
    "    # Load image.\n",
    "    img = PIL.Image.open(src_file)\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, PIL.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 = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB')\n",
    "        quad += pad[:2]\n",
    "        \n",
    "    # Transform.\n",
    "    img = img.transform((transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR)\n",
    "    if output_size < transform_size:\n",
    "        img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS)\n",
    "     \n",
    "     \n",
    "    #display(img)\n",
    "    # Save aligned image.\n",
    "    img.save(dst_file, quality=100, subsampling=0)\n",
    "\n",
    "landmarks_detector = face_alignment.FaceAlignment(face_alignment.LandmarksType._3D, flip_input=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e985e9e-5112-4347-8fe1-81232362d2c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#get face image from frame\n",
    "for frameFile in tqdm(glob(frameDir+\"/*.jpg\")):\n",
    "    name=Path(frameFile).stem\n",
    "\n",
    "    \n",
    "    ######################## use anime face detector landmark to align\n",
    "    \n",
    "    image = cv2.imread(frameFile)\n",
    "    preds = detector(image)\n",
    "    for i, face_landmark in enumerate(preds):\n",
    "        if face_landmark[\"bbox\"][4]<0.5 or np.mean(face_landmark[\"keypoints\"][:,2])<0.3: continue #skip low confidence\n",
    "\n",
    "        aligned_face_path = os.path.join(alignedDir, name+\"_\"+str(i).zfill(4)+\".jpg\")\n",
    "        image_align(frameFile, aligned_face_path, face_landmark[\"keypoints\"],use_landmark_28=True)\n",
    "\n",
    "    ######################## use face-alignment landmark to align\n",
    "    \n",
    "#     face_landmarks=landmarks_detector.get_landmarks(frameFile)\n",
    "#     if face_landmarks is None: \n",
    "#         continue  #skip none output\n",
    "#     for i, face_landmark in enumerate(face_landmarks):\n",
    "#         aligned_face_path = os.path.join(alignedDir, name+\"_\"+str(i).zfill(4)+\".jpg\")\n",
    "#         image_align(frameFile, aligned_face_path, face_landmark)\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e90193c9-f23f-4244-b55f-f1942ef4f7ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "#filter blurry image\n",
    "threshold=70\n",
    "\n",
    "for i,file in tqdm(enumerate(glob(alignedDir+\"/*.jpg\"))):\n",
    "    name=Path(file).name\n",
    "    image = cv2.imread(file)\n",
    "    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "    # fm= cv2.Laplacian(gray, cv2.CV_64F).var()\n",
    "    fm=np.max(cv2.convertScaleAbs(cv2.Laplacian(gray,3)))\n",
    "    \n",
    "    if threshold < fm:        \n",
    "        # display(Image.open(file))\n",
    "        copyfile(file, filteredDir+\"/\"+name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37ab946e-1323-4135-bfaa-6389c1e4928d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# color correction and denoise\n",
    "def better_cb(img, percent=1):\n",
    "    # from https://github.com/luftj/MaRE/blob/4284fe2b3307ca407e87e3b0dbdaa3c1ef646731/simple_cb.py\n",
    "\n",
    "    if not percent or percent == 0 or percent == 100:\n",
    "        return img\n",
    "\n",
    "    out_channels = []\n",
    "    cumstops = (\n",
    "        img.shape[0] * img.shape[1] * percent / 200.0,\n",
    "        img.shape[0] * img.shape[1] * (1 - percent / 200.0),\n",
    "    )\n",
    "    for channel in cv2.split(img):\n",
    "        cumhist = np.cumsum(cv2.calcHist([channel], [0], None, [256], (0, 256)))\n",
    "        low_cut, high_cut = np.searchsorted(cumhist, cumstops)\n",
    "        lut = np.concatenate(\n",
    "            (\n",
    "                np.zeros(low_cut),\n",
    "                np.around(np.linspace(0, 255, high_cut - low_cut + 1)),\n",
    "                255 * np.ones(255 - high_cut),\n",
    "            )\n",
    "        )\n",
    "        out_channels.append(cv2.LUT(channel, lut.astype(\"uint8\")))\n",
    "    return cv2.merge(out_channels)\n",
    "\n",
    "\n",
    "for i, file in tqdm(enumerate(glob(filteredDir + \"/*.jpg\"))):\n",
    "    name = Path(file).name\n",
    "    image = cv2.imread(file)\n",
    "    image = better_cb(image)  # color correction\n",
    "    image = cv2.fastNlMeansDenoisingColored(image, None, 3, 3, 7, 21)  # denoise\n",
    "    img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n",
    "    img.save(preprocessedDir + \"/\" + name, quality=100, subsampling=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f5e707d-6f57-4f01-8127-37adc321d479",
   "metadata": {},
   "outputs": [],
   "source": [
    "#display\n",
    "for i,file in tqdm(enumerate(glob(preprocesseddDir+\"/*.jpg\")[0:10])):\n",
    "    display(Image.open(file))\n",
    "print(len(glob(filteredDir+\"/*.jpg\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00e960ea-dae5-46a0-abdd-44ca8e365492",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make zip file\n",
    "!cd stylegan3 && python dataset_tool.py --source=\"../$preprocessedDir\" --dest=\"../$dataZip\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68351e85-8e38-4247-92ec-ffd8f2e8bf10",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tf2]",
   "language": "python",
   "name": "conda-env-tf2-py"
  },
  "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}