unet_generate_masks.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the necessary packages\n",
    "from imutils import paths \n",
    "import argparse\n",
    "import imutils\n",
    "import json \n",
    "import cv2\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "# tools for parallelism\n",
    "from multiprocessing import Pool\n",
    "from functools import partial\n",
    "\n",
    "# separate function to compute mask for one landmark\n",
    "from utils.utils import one_key_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# original image size and intended image size\n",
    "H = 348\n",
    "W = 464\n",
    "\n",
    "h = 256\n",
    "w = 256\n",
    "\n",
    "# number of facial landmarks\n",
    "KEYPOINTS = 54"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_masks(jsonFolder):\n",
    "    challenging_imgs_csv = \"dataset/challenging_images.csv\"\n",
    "    \n",
    "    # folder where generated masks will be stored\n",
    "    masksFolder = jsonFolder.replace('/json', '/masks')\n",
    "    if not os.path.exists(masksFolder):\n",
    "        os.makedirs(masksFolder)\n",
    "\n",
    "    # paths to json files\n",
    "    jsonPaths = list(paths.list_files(jsonFolder, validExts=\"json\"))\n",
    "    jsonPaths = sorted(jsonPaths)\n",
    "\n",
    "    # loop over the json files\n",
    "    for ind, jsonPath in enumerate(jsonPaths, 1):\n",
    "\n",
    "        print(\"[INFO] Processing {} file ({}/{})\".format(jsonPath.split(\"/\")[-1], ind, len(jsonPaths)))\n",
    "\n",
    "        # opening the json file \n",
    "        f = open(jsonPath,) \n",
    "\n",
    "        # returns the json object as a dictionary \n",
    "        data = json.load(f) \n",
    "\n",
    "        # extracting filename\n",
    "        filename = jsonPath.split(\"/\")[-1]\n",
    "        filename = filename.split(\".json\")[0]\n",
    "\n",
    "        # extracting trial data from filename\n",
    "        sub_id, trial_id, exp_id, pos_id, *rest = filename.split(\"_\")\n",
    "        info = \"{}_{}_{}\".format(sub_id, trial_id, pos_id)\n",
    "\n",
    "        # if image is challenging, skip it\n",
    "        df = pd.read_csv(challenging_imgs_csv)\n",
    "\n",
    "        if ((df['sub_id'] == int(sub_id)) & (df['trial_id'] == int(trial_id)) & (df['pos_id'] == int(pos_id))).any():\n",
    "            print('Excluding: ', info)\n",
    "            continue\n",
    "\n",
    "        # iterating through the shapes\n",
    "        landmarks = []\n",
    "        for shape in data['shapes']:\n",
    "            # face label corresponds to bounding boxes\n",
    "            if shape['label'] == 'face':\n",
    "                [[xs, ys], [xe, ye]] = shape['points']\n",
    "                (xs, ys, xe, ye) = (int(xs), int(ys), int(xe), int(ye))\n",
    "\n",
    "            # other labels correspond to landmarks\n",
    "            elif shape['label'] != 'face':\n",
    "                # iterating through the landmarks\n",
    "                for point in shape['points']:\n",
    "                    (x, y) = point\n",
    "                    (x, y) = (int(x), int(y))\n",
    "\n",
    "                    # adjust the landmark positions w.r.t. bounding boxes and new image size\n",
    "                    (crop_x, crop_y) = (x-xs, y-ys)\n",
    "                    (crop_x, crop_y) = (crop_x / (xe-xs) * w, crop_y / (ye-ys) * h)\n",
    "                    landmarks.append([crop_y, crop_x])\n",
    "\n",
    "        landmarks = np.array(landmarks)\n",
    "\n",
    "        # array to contain the masks for each point\n",
    "        B = np.zeros((KEYPOINTS, h, w))\n",
    "\n",
    "        # mask generation is computed on CPU\n",
    "        # hence, the landmark computation is manually divided to 6 processors\n",
    "        # separate landmarks to 6 subsets for parallel computation\n",
    "        keys1 = (list(range(0, 9)))\n",
    "        keys2 = (list(range(9, 18)))\n",
    "        keys3 = (list(range(18, 27)))\n",
    "        keys4 = (list(range(27, 36)))\n",
    "        keys5 = (list(range(36, 45)))\n",
    "        keys6 = (list(range(45, 54)))\n",
    "\n",
    "        with Pool(processes=6) as pool:\n",
    "            b1 = pool.map(partial(one_key_mask, landmarks = landmarks, m=h, n=w), keys1)\n",
    "            b2 = pool.map(partial(one_key_mask, landmarks = landmarks, m=h, n=w), keys2)\n",
    "            b3 = pool.map(partial(one_key_mask, landmarks = landmarks, m=h, n=w), keys3)\n",
    "            b4 = pool.map(partial(one_key_mask, landmarks = landmarks, m=h, n=w), keys4)\n",
    "            b5 = pool.map(partial(one_key_mask, landmarks = landmarks, m=h, n=w), keys5)\n",
    "            b6 = pool.map(partial(one_key_mask, landmarks = landmarks, m=h, n=w), keys6)\n",
    "\n",
    "            B[0:9, :, :] = np.array(b1)\n",
    "            B[9:18, :, :] = np.array(b2)\n",
    "            B[18:27, :, :] = np.array(b3)\n",
    "            B[27:36, :, :] = np.array(b4)\n",
    "            B[36:45, :, :] = np.array(b5)\n",
    "            B[45:54, :, :] = np.array(b6)\n",
    "\n",
    "        # combining to one mask\n",
    "        mask = np.sum(B, axis=0)\n",
    "        mask = np.clip(mask, 0, 1)\n",
    "        mask = mask * 255.0\n",
    "        mask = mask.astype('uint8')\n",
    "\n",
    "        # saving the mask\n",
    "        maskPath = os.path.join(masksFolder, \"{}.png\".format(filename))\n",
    "        cv2.imwrite(maskPath, mask)\n",
    "\n",
    "        # Closing file \n",
    "        f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_masks('dataset/gray/train/json')\n",
    "generate_masks('dataset/gray/val/json')\n",
    "generate_masks('dataset/gray/test/json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}