training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Facial Recognition Training Pipeline Example\n",
    "This notebook takes a directory of training images for a number of different people and trains a classifier to do facial recognition. The method is as follows:\n",
    "    1. Crop and format input images\n",
    "    2. Import model\n",
    "    3. Perform transfer learning using fast.ai\n",
    "    4. Output in ONNX format for use\n",
    "    \n",
    "Additionally, in order to load training images from their original directory, it must be structured as below:\n",
    "\n",
    "<code>raw_training_data_dir <br>│<br>└───person_1<br>│ │───IMG1<br>│ │───IMG2<br>│ │ ....<br>└───person_2<br>| │───IMG1<br>| │───IMG2<br>| | ....</code>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from fastai import *\n",
    "from fastai.vision import *\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.dataset import TensorDataset\n",
    "import torch.utils.model_zoo as model_zoo\n",
    "import torch.onnx\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2, dlib\n",
    "import os, shutil, io, random, subprocess\n",
    "os.system('pip install facenet_pytorch')\n",
    "from facenet_pytorch import MTCNN, InceptionResnetV1\n",
    "from datetime import datetime\n",
    "\n",
    "bs=8\n",
    "torch.cuda.is_available()\n",
    "torch.cuda.set_device(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Functions utilizing OpenCV for image loading and manipulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to create list of cv2-loaded images in a given directory\n",
    "def load_images_from_folder(folder):\n",
    "    images = []\n",
    "    for filename in os.listdir(folder):\n",
    "        img = cv2.imread(os.path.join(folder,filename),1)\n",
    "        if img is not None:\n",
    "            images.append(img)\n",
    "    return images\n",
    "\n",
    "# function to convert cv2-loaded image into RGB\n",
    "def convertToRGB(image):\n",
    "    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "# scale image down or up while keeping original aspect ratio\n",
    "def scale_image(image, ratio):\n",
    "    scale_percent = ratio # percent of original size -> change to specify new size\n",
    "    width = int(img.shape[1] * scale_percent / 100)\n",
    "    height = int(img.shape[0] * scale_percent / 100)\n",
    "    dim = (width, height)\n",
    "    # resize image\n",
    "    return cv2.resize(image, dim, interpolation = cv2.INTER_AREA)\n",
    "\n",
    "# returns all faces detected in input image using a DNN\n",
    "def detectFaceOpenCVDnn(net, image):\n",
    "    frameHeight = image.shape[0]\n",
    "    frameWidth = image.shape[1]\n",
    "    \n",
    "    data = cv2.dnn.blobFromImage(image, 1.0, (300, 300), [104, 117, 123], False, False)\n",
    "    \n",
    "    net.setInput(data)\n",
    "    detections = net.forward()\n",
    "    bounding_boxes = []\n",
    "    for i in range(detections.shape[2]):\n",
    "        confidence = detections[0, 0, i, 2]\n",
    "        # print(confidence)\n",
    "        if confidence > conf_threshold:\n",
    "            x1 = int(detections[0, 0, i, 3] * frameWidth)\n",
    "            y1 = int(detections[0, 0, i, 4] * frameHeight)\n",
    "            x2 = int(detections[0, 0, i, 5] * frameWidth)\n",
    "            y2 = int(detections[0, 0, i, 6] * frameHeight)\n",
    "            if not x1 < 0 and not x2 > frameWidth and not y1 < 0 and not y2 > frameHeight:\n",
    "                bounding_boxes.append([x1, y1, x2, y2])\n",
    "                # cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), int(round(frameHeight/150)), 8)\n",
    "    return bounding_boxes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preprocessing of input images \n",
    "    1. Uses pretrained Caffe or TensorFlow model to perform face detection\n",
    "    2. Crops training photos around detected faces\n",
    "    3. Save cropped photos in new training set with same structure for loading into DataBunch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# optional haar model to use to locate faces\n",
    "# haar_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')\n",
    "\n",
    "# will only choose partitions of the image that the DNN says minimum conf_threshold % probability of being a face\n",
    "conf_threshold = 0.98\n",
    "\n",
    "# directory containing images of people to train on. NOTE: must be in same directory as this notebook\n",
    "raw_training_data_dir = 'intern_images'\n",
    "\n",
    "#load DNN models\n",
    "DNN = \"CAFFE\"\n",
    "if DNN == \"CAFFE\":\n",
    "    modelFile = \"models/weights.caffemodel\"\n",
    "    configFile = \"models/deploy.prototxt\"\n",
    "    net = cv2.dnn.readNetFromCaffe(configFile, modelFile)\n",
    "else:\n",
    "    modelFile = \"opencv_face_detector_uint8.pb\"\n",
    "    configFile = \"opencv_face_detector.pbtxt\"\n",
    "    net = cv2.dnn.readNetFromTensorflow(modelFile, configFile)\n",
    "\n",
    "# make directory for training images, set readable\n",
    "if not os.path.exists('training_set'):\n",
    "    os.mkdir('training_set', 0o777)\n",
    "elif os.path.isdir('training_set'): \n",
    "    shutil.rmtree('training_set')\n",
    "    os.mkdir('training_set', 0o777)\n",
    "\n",
    "# load all input images, process and output cropped faces to new test directory\n",
    "for dirname, dirnames, filenames in os.walk('./' + raw_training_data_dir):\n",
    "    for subdirname in dirnames:\n",
    "        \n",
    "        # create subpath for specific class images for after cropping, and delete existing images if they exist\n",
    "        sub_path = 'training_set/' + subdirname\n",
    "        if not os.path.exists(sub_path):\n",
    "            os.mkdir(sub_path, 0o777)\n",
    "        elif os.path.isdir(sub_path): \n",
    "            shutil.rmtree(sub_path)\n",
    "            os.mkdir(sub_path, 0o777)\n",
    "            \n",
    "        \n",
    "        #print(os.path.join(dirname, subdirname))\n",
    "        imgs = load_images_from_folder(os.path.join(dirname, subdirname))\n",
    "        # print('loaded images for '+ subdirname)\n",
    "        index = 0\n",
    "        for img in imgs:\n",
    "            # uncomment to \n",
    "            # faces = haar_cascade.detectMultiScale(img, scaleFactor = 1.2, minNeighbors = 4, minSize=(500, 500))\n",
    "            \n",
    "            # scale image down and detect faces using the DNN\n",
    "            img = scale_image(img, 18)\n",
    "            faces = detectFaceOpenCVDnn(net,img)\n",
    "            # if there is one face in the frame, crop the frame to isolate the face\n",
    "            if len(faces) >= 1:\n",
    "                index+=1\n",
    "                (x1,y1,x2,y2) = faces[0]\n",
    "                cropped = img[y1:y2, x1:x2].copy()\n",
    "                # uncomment to see all loaded pictures\n",
    "                plt.figure()\n",
    "                plt.imshow(convertToRGB(cropped))\n",
    "                \n",
    "                # create and store cropped face images\n",
    "                filename = subdirname + \"_\" + '{0:04d}'.format(index)+\".jpg\"\n",
    "                cv2.imwrite(os.path.join(sub_path , filename), cropped)\n",
    "        print(\"loaded \" + str(len(imgs)) + \" images of \" + subdirname)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load cropped test faces into ImageDataBunch and apply transforms\n",
    "    Transforms include jitter, brightness adjustment, and contrast adjustment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## applying transforms\n",
    "tfms_list = [jitter(magnitude=(random.randrange(-3,3)/100), p=0.25), \n",
    "             contrast(scale=(0.5, 2.), p=0.5), brightness(change=(0.1, 0.9), p=0.5)]\n",
    "#tfms_list = [rotate(degrees=270, p=1)]\n",
    "tfms = [tfms_list, tfms_list]\n",
    "\n",
    "## Declaring path of dataset\n",
    "path_img = Path('/root/model_training/training_set')\n",
    "\n",
    "## Loading data \n",
    "data = ImageDataBunch.from_folder(path=path_img, train='/', valid_pct=0.07, ds_tfms=tfms, bs=bs, size=(150,150))\n",
    "\n",
    "## Normalizing data based on ImageNet parameters\n",
    "data.normalize(imagenet_stats)\n",
    "\n",
    "## Showing some contents of databunch to confirm it is loaded correctly with transforms\n",
    "data.show_batch(rows=2)\n",
    "print(data.classes)\n",
    "len(data.classes),data.c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "### Transfer learning on pretrained facenet_pytorch model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_model(pretrained=True):\n",
    "    if pretrained:\n",
    "        facenet_base = InceptionResnetV1(pretrained='vggface2').eval()\n",
    "    else:\n",
    "        facenet_base = models.resnet18\n",
    "    return facenet_base\n",
    "\n",
    "## To create a model with pretrained weights\n",
    "learn = cnn_learner(data, get_model, metrics=accuracy)\n",
    "# learn.freeze_to(300)\n",
    "learn.fit_one_cycle(10,1e-2)\n",
    "learn.save(\"face_detection_\" + str(datetime.utcnow().strftime(\"%m%d%Y\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "print('Training results:  ' + str(learn.validate(learn.data.train_dl)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('Validation results: ' + str(learn.validate(learn.data.valid_dl)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(learn.loss_func)\n",
    "preds,y,losses = learn.get_preds(with_loss=True)\n",
    "interp = ClassificationInterpretation(learn, preds, y, losses)\n",
    "interp = ClassificationInterpretation.from_learner(learn)\n",
    "interp.plot_top_losses(9, figsize=(10,10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "interp.plot_confusion_matrix(figsize=(5,5), dpi=60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot()\n",
    "plt.plot(learn.recorder.lrs[10:-5],learn.recorder.losses[10:-5],'r.')\n",
    "plt.xscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Load pretrained model weights\n",
    "for file in os.listdir(\"/root/model_training/training_set/models\"):\n",
    "    if file.endswith(\".pth\"):\n",
    "        model_load = os.path.join(\"/mydir\", file)\n",
    "batch_size = 16    # just a random number\n",
    "\n",
    "# Initialize model with the pretrained weights\n",
    "map_location = lambda storage, loc: storage\n",
    "if torch.cuda.is_available():\n",
    "    map_location = None\n",
    "# learn.load_state_dict(model_zoo.load_url(model_load, map_location=map_location))\n",
    "\n",
    "torch_model = learn.model\n",
    "\n",
    "# set the train mode to false since we will only run the forward pass.\n",
    "torch_model.train(False)\n",
    "\n",
    "# Input to the model\n",
    "x = torch.randn(10, 3, 224, 224, device='cuda',requires_grad=True)\n",
    "\n",
    "view_output=False\n",
    "\n",
    "# Export the model\n",
    "torch_out = torch.onnx._export(torch_model,             # model being run\n",
    "                               x,                       # model input (or a tuple for multiple inputs)\n",
    "                               \"models/face_detector.onnx\",    # where to save the model (can be a file or file-like object)\n",
    "                               export_params=True,      # store the trained parameter weights inside the model file\n",
    "                               verbose=view_output)     # whether or not to view onncxoutput progress\n",
    "os.chmod(\"models/face_detector.onnx\", 0o777)\n",
    "print(\"loaded model in onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only run if you want to use caffe2 model in your application\n",
    "\n",
    "os.system(\"pip install onnx\")\n",
    "import onnx\n",
    "import caffe2.python.onnx.backend as backend\n",
    "\n",
    "# Load the ONNX ModelProto object. model is a standard Python protobuf object\n",
    "model = onnx.load(\"face_detector.onnx\")\n",
    "\n",
    "# Check that the IR is well formed\n",
    "onnx.checker.check_model(model)\n",
    "\n",
    "rep = backend.prepare(model, device=\"CUDA:0\")\n",
    "\n",
    "outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))\n",
    "\n",
    "os.system(\"python convertCaffe.py ./models/face_detector.onnx + ./models/face_detector_\"\n",
    "        +str(datetime.utcnow().strftime(\"%m%d%Y\"))+\".prototxt ./models/face_detector_\"\n",
    "        +str(datetime.utcnow().strftime(\"%m%d%Y\"))+\".caffemodel\")"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}