face_tracking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Face tracking pipeline\n",
    "\n",
    "The following example illustrates how to use the `facenet_pytorch` python package to perform face detection and tracking on an image dataset using MTCNN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from facenet_pytorch import MTCNN\n",
    "import torch\n",
    "import numpy as np\n",
    "import mmcv, cv2\n",
    "from PIL import Image, ImageDraw\n",
    "from IPython import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Determine if an nvidia GPU is available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on device: cpu\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/t/Desktop/raschka/lib/python3.11/site-packages/torch/cuda/__init__.py:138: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 11040). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:108.)\n",
      "  return torch._C._cuda_getDeviceCount() > 0\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "print('Running on device: {}'.format(device))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define MTCNN module\n",
    "\n",
    "Note that, since MTCNN is a collection of neural nets and other code, the device must be passed in the following way to enable copying of objects when needed internally.\n",
    "\n",
    "See `help(MTCNN)` for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "mtcnn = MTCNN(keep_all=True, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Get a sample video\n",
    "\n",
    "We begin by loading a video with some faces in it. The `mmcv` PyPI package by mmlabs is used to read the video frames (it can be installed with `pip install mmcv`). Frames are then converted to PIL images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'mmcv' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m video \u001b[38;5;241m=\u001b[39m \u001b[43mmmcv\u001b[49m\u001b[38;5;241m.\u001b[39mVideoReader(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvideo.mp4\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m      2\u001b[0m frames \u001b[38;5;241m=\u001b[39m [Image\u001b[38;5;241m.\u001b[39mfromarray(cv2\u001b[38;5;241m.\u001b[39mcvtColor(frame, cv2\u001b[38;5;241m.\u001b[39mCOLOR_BGR2RGB)) \u001b[38;5;28;01mfor\u001b[39;00m frame \u001b[38;5;129;01min\u001b[39;00m video]\n\u001b[1;32m      4\u001b[0m display\u001b[38;5;241m.\u001b[39mVideo(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvideo.mp4\u001b[39m\u001b[38;5;124m'\u001b[39m, width\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m640\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'mmcv' is not defined"
     ]
    }
   ],
   "source": [
    "video = mmcv.VideoReader('video.mp4')\n",
    "frames = [Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) for frame in video]\n",
    "\n",
    "display.Video('video.mp4', width=640)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run video through MTCNN\n",
    "\n",
    "We iterate through each frame, detect faces, and draw their bounding boxes on the video frames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames_tracked = []\n",
    "for i, frame in enumerate(frames):\n",
    "    print('\\rTracking frame: {}'.format(i + 1), end='')\n",
    "    \n",
    "    # Detect faces\n",
    "    boxes, _ = mtcnn.detect(frame)\n",
    "    \n",
    "    # Draw faces\n",
    "    frame_draw = frame.copy()\n",
    "    draw = ImageDraw.Draw(frame_draw)\n",
    "    for box in boxes:\n",
    "        draw.rectangle(box.tolist(), outline=(255, 0, 0), width=6)\n",
    "    \n",
    "    # Add to frame list\n",
    "    frames_tracked.append(frame_draw.resize((640, 360), Image.BILINEAR))\n",
    "print('\\nDone')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Display detections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = display.display(frames_tracked[0], display_id=True)\n",
    "i = 1\n",
    "try:\n",
    "    while True:\n",
    "        d.update(frames_tracked[i % len(frames_tracked)])\n",
    "        i += 1\n",
    "except KeyboardInterrupt:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save tracked video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"FR\")\n",
    "dim = frames_tracked[0].size\n",
    "print(dim)\n",
    "fourcc = cv2.VideoWriter_fourcc(*'FMP4')    \n",
    "print(fourcc)\n",
    "video_tracked = cv2.VideoWriter('./video_tracked.mkv', fourcc, 25.0, dim)\n",
    "print(frames_tracked)\n",
    "for frame in frames_tracked:\n",
    "    video_tracked.write(cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR))\n",
    "video_tracked.release()\n",
    "print('done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}