ffhq-align-faces.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3iodGlkx3fc9"
   },
   "source": [
    "# Align Faces\n",
    "by @duskvirkus\n",
    "Thanks to Derrick Schultz for the orginal notebook.\n",
    "\n",
    "This notebook will walk thru the process to take a folder of face images and convert them to images in the format of FFHQ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# notebook utils\n",
    "import os\n",
    "root_dir = os.getcwd()\n",
    "\n",
    "vendor_path = os.path.join(root_dir, 'vendor')\n",
    "nbt_path = os.path.join(vendor_path, 'notebook-utils');\n",
    "    \n",
    "if not os.path.exists(nbt_path):\n",
    "    if not os.path.exists(vendor_path):\n",
    "        os.makedirs(vendor_path)\n",
    "    %cd {vendor_path}\n",
    "    !git clone https://github.com/duskvirkus/notebook-utils\n",
    "    %cd {root_dir}\n",
    "else:\n",
    "    %cd {nbt_path}\n",
    "    !git pull origin main\n",
    "    %cd {root_dir}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_folder = '/home/duskvirkus/dataset-creation/painterly-faces/raw/bj00100'\n",
    "output_folder = '/home/duskvirkus/dataset-creation/painterly-faces/faces/bj00100'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.insert(1, os.path.join(os.getcwd(), 'vendor/notebook-utils'))\n",
    "from user_input import yes_or_no\n",
    "from progress import update_progress\n",
    "import shutil\n",
    "\n",
    "%cd data/\n",
    "\n",
    "dir_name = 'align-faces'\n",
    "\n",
    "if not os.path.exists(dir_name):\n",
    "    os.makedirs(dir_name)\n",
    "else:\n",
    "    if yes_or_no('Delete existing working directory and start fresh?'):\n",
    "        shutil.rmtree(dir_name)\n",
    "        os.makedirs(dir_name)\n",
    "        \n",
    "%cd {dir_name}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EgnXI-bk3ZIt"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists('stylegan2-ada'):\n",
    "    !git clone https://github.com/duskvirkus/stylegan2-ada\n",
    "%cd stylegan2-ada\n",
    "working_dir = os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "52h5HDcq33oh"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import bz2\n",
    "from tensorflow.keras.utils import get_file\n",
    "from ffhq_dataset.face_alignment import image_align\n",
    "from ffhq_dataset.landmarks_detector import LandmarksDetector\n",
    "\n",
    "LANDMARKS_MODEL_URL = 'http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2'\n",
    "RAW_IMAGES_DIR = input_folder\n",
    "ALIGNED_IMAGES_DIR = output_folder\n",
    "\n",
    "if not os.path.exists(ALIGNED_IMAGES_DIR):\n",
    "        os.makedirs(ALIGNED_IMAGES_DIR)\n",
    "\n",
    "\n",
    "def unpack_bz2(src_path):\n",
    "    data = bz2.BZ2File(src_path).read()\n",
    "    dst_path = src_path[:-4]\n",
    "    with open(dst_path, 'wb') as fp:\n",
    "        fp.write(data)\n",
    "    return dst_path\n",
    "\n",
    "landmarks_model_path = unpack_bz2(get_file('shape_predictor_68_face_landmarks.dat.bz2', LANDMARKS_MODEL_URL, cache_subdir='temp'))\n",
    "\n",
    "\n",
    "landmarks_detector = LandmarksDetector(landmarks_model_path)\n",
    "\n",
    "full_count = len([x for x in os.listdir(RAW_IMAGES_DIR) if x[0] not in '._'])\n",
    "count = 0\n",
    "\n",
    "for img_name in [x for x in os.listdir(RAW_IMAGES_DIR) if x[0] not in '._']:\n",
    "    raw_img_path = os.path.join(RAW_IMAGES_DIR, img_name)\n",
    "    for i, face_landmarks in enumerate(landmarks_detector.get_landmarks(raw_img_path), start=1):\n",
    "        face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], i)\n",
    "        aligned_face_path = os.path.join(ALIGNED_IMAGES_DIR, face_img_name)\n",
    "        os.makedirs(ALIGNED_IMAGES_DIR, exist_ok=True)\n",
    "        image_align(raw_img_path, aligned_face_path, face_landmarks, x_scale=1.2, em_scale=0.5)\n",
    "    count += 1\n",
    "    update_progress('align faces images progress', count / full_count)\n",
    "\n",
    "update_progress('align faces images progress', 1)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Copy of Align Faces.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "jupyter-ml-art",
   "language": "python",
   "name": "jupyter-ml-art"
  },
  "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}