Social-Distance-Tool.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Social Distance Tool with depth\n",
    "\n",
    "This tool combines two algorithms to accurately detect people who are violating the social distancing protocol:\n",
    "- Facebook/Detectron2 (Faster RCNN implementation)`https://github.com/facebookresearch/detectron2`\n",
    "- \"Digging into Self-Supervised Monocular Depth Prediction\" `https://github.com/nianticlabs/monodepth2`\n",
    "\n",
    "**Input:**\n",
    "- A video sequence\n",
    "\n",
    "**Output:**\n",
    "- bounding boxes on all persons detected in the video\n",
    "- highlighing people who are in close proximity\n",
    "- depth map for accurate calculations \n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code\n",
    "**Import libraries for Detectron2**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python -m detectron2.utils.collect_env # to check if Detectron2 is working fine\n",
    "# Some basic setup:\n",
    "# Setup detectron2 logger\n",
    "import detectron2\n",
    "from detectron2.utils.logger import setup_logger\n",
    "setup_logger()\n",
    "\n",
    "# import some common libraries\n",
    "import numpy as np\n",
    "import cv2\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# import some common detectron2 utilities\n",
    "from detectron2 import model_zoo\n",
    "from detectron2.engine import DefaultPredictor\n",
    "from detectron2.config import get_cfg\n",
    "from detectron2.utils.visualizer import Visualizer\n",
    "from detectron2.data import MetadataCatalog"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Import libraries and files for MonoDepth2 algorithm**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Libraries for monodepth2\n",
    "from __future__ import absolute_import, division, print_function\n",
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "import PIL.Image as pil\n",
    "\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "\n",
    "import networks\n",
    "from utils import download_model_if_doesnt_exist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Convert Video to PNG Frames**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rm: cannot remove 'frames/*': No such file or directory\n",
      "CPU times: user 8.19 s, sys: 144 ms, total: 8.33 s\n",
      "Wall time: 8.17 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "!rm -r frames/*\n",
    "!mkdir frames/\n",
    "\n",
    "#specify path to video\n",
    "video = \"sample.mp4\"\n",
    "\n",
    "#capture video\n",
    "cap = cv2.VideoCapture(video)\n",
    "cnt=0\n",
    "FPS=cap.get(cv2.CAP_PROP_FPS)\n",
    "# Check if video file is opened successfully\n",
    "if (cap.isOpened()== False): \n",
    "  print(\"Error opening video stream or file\")\n",
    "\n",
    "ret,first_frame = cap.read()\n",
    "\n",
    "#Read until video is completed\n",
    "while(cap.isOpened()):\n",
    "\n",
    "  # Capture frame-by-frame\n",
    "  ret, frame = cap.read()\n",
    "\n",
    "  if ret == True:\n",
    "\n",
    "    #save each frame to folder        \n",
    "    cv2.imwrite('frames/'+'{:04d}'.format(cnt)+'.png', frame)\n",
    "    cnt=cnt+1\n",
    "    if(cnt==150):\n",
    "      break\n",
    "\n",
    "  # Break the loop\n",
    "  else: \n",
    "    break\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Download a pretrained model from Detectron2 Model Zoo**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg = get_cfg()\n",
    "\n",
    "# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library\n",
    "cfg.merge_from_file(model_zoo.get_config_file(\"COCO-Detection/faster_rcnn_R_50_C4_3x.yaml\"))\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9  # set threshold for this model\n",
    "\n",
    "# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well\n",
    "cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(\"COCO-Detection/faster_rcnn_R_50_C4_3x.yaml\")\n",
    "predictor = DefaultPredictor(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define all the key functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a function which return the bottom center of every bbox\n",
    "def mid_point(img,person,idx):\n",
    "  #get the coordinates\n",
    "  x1,y1,x2,y2 = person[idx]\n",
    "  _ = cv2.rectangle(img, (x1, y1), (x2, y2), (0,0,255), 2)\n",
    "  \n",
    "  #compute bottom center of bbox\n",
    "  x_mid = int((x1+x2)/2)\n",
    "  y_mid = int(y2)\n",
    "  mid   = (x_mid,y_mid)\n",
    "  \n",
    "  _ = cv2.circle(img, mid, 5, (0, 0, 255), -1)\n",
    "  cv2.putText(img, str(idx), mid, cv2.FONT_HERSHEY_SIMPLEX,1, (255, 255, 255), 2, cv2.LINE_AA)\n",
    "  \n",
    "  return mid\n",
    "\n",
    "# define a function which computes euclidean distance between two midpoints\n",
    "from scipy.spatial import distance\n",
    "def compute_distance(midpoints,num):\n",
    "  dist = np.zeros((num,num))\n",
    "  for i in range(num):\n",
    "    for j in range(i+1,num):\n",
    "      if i!=j:\n",
    "        dst = distance.euclidean(midpoints[i], midpoints[j])\n",
    "        dist[i][j]=dst\n",
    "  return dist\n",
    "\n",
    "\n",
    "# Finds pairs of people who are close together\n",
    "def find_closest(dist,num,thresh):\n",
    "  p1=[]\n",
    "  p2=[]\n",
    "  d=[]\n",
    "  for i in range(num):\n",
    "    for j in range(i,num):\n",
    "      if( (i!=j) & (dist[i][j]<=thresh)):\n",
    "        p1.append(i)\n",
    "        p2.append(j)\n",
    "        d.append(dist[i][j])\n",
    "  return p1,p2,d\n",
    "\n",
    "\n",
    "# Given pairs of people who are close, color them red\n",
    "def change_2_red(img,person,p1,p2):\n",
    "  mid1 = []\n",
    "  mid2 = []\n",
    "  for p in p1:\n",
    "    mid1.append(mid_point(img,person,p))\n",
    "  for pp in p2:\n",
    "    mid2.append(mid_point(img,person,pp))\n",
    "  for inx in range(len(mid1)):\n",
    "      _ = cv2.line(img, mid1[inx], mid2[inx], (0,255,0), thickness=2, lineType=8, shift=0)\n",
    "  \n",
    "  risky = np.unique(p1+p2)\n",
    "  for i in risky:\n",
    "    x1,y1,x2,y2 = person[i]\n",
    "    _ = cv2.rectangle(img, (x1, y1), (x2, y2), (255,0,0), 2)  \n",
    "  return img\n",
    "\n",
    "\n",
    "# Main function to find closest people\n",
    "def find_closest_people(name,thresh,savedir):\n",
    "\n",
    "  img = cv2.imread('frames/'+name)\n",
    "  outputs = predictor(img)\n",
    "  classes=outputs['instances'].pred_classes.cpu().numpy()\n",
    "  bbox=outputs['instances'].pred_boxes.tensor.cpu().numpy()\n",
    "  ind = np.where(classes==0)[0]\n",
    "  person=bbox[ind]\n",
    "  midpoints = [mid_point(img,person,i) for i in range(len(person))]\n",
    "  num = len(midpoints)\n",
    "  dist= compute_distance(midpoints,num)\n",
    "  p1,p2,d=find_closest(dist,num,thresh)\n",
    "  img = change_2_red(img,person,p1,p2)\n",
    "  cv2.imwrite(savedir+'/'+name,img)\n",
    "  return 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Fetch all the frames of the video sequence**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames=[]\n",
    "for file in os.listdir(\"frames/\"):\n",
    "    if file.endswith(\".png\"):\n",
    "        frames.append(file)\n",
    "frames.sort()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Main loop to get all the files**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  0%|          | 0/150 [00:00<?, ?it/s]\u001b[A\n",
      "  1%|          | 1/150 [00:00<01:23,  1.79it/s]\u001b[A\n",
      "  1%|▏         | 2/150 [00:00<01:15,  1.97it/s]\u001b[A\n",
      "  2%|▏         | 3/150 [00:01<01:09,  2.12it/s]\u001b[A\n",
      "  3%|▎         | 4/150 [00:01<01:05,  2.24it/s]\u001b[A\n",
      "  3%|▎         | 5/150 [00:02<01:02,  2.33it/s]\u001b[A\n",
      "  4%|▍         | 6/150 [00:02<01:00,  2.40it/s]\u001b[A\n",
      "  5%|▍         | 7/150 [00:02<00:58,  2.45it/s]\u001b[A\n",
      "  5%|▌         | 8/150 [00:03<00:57,  2.49it/s]\u001b[A\n",
      "  6%|▌         | 9/150 [00:03<00:56,  2.51it/s]\u001b[A\n",
      "  7%|▋         | 10/150 [00:04<00:55,  2.52it/s]\u001b[A\n",
      "  7%|▋         | 11/150 [00:04<00:54,  2.54it/s]\u001b[A\n",
      "  8%|▊         | 12/150 [00:04<00:54,  2.55it/s]\u001b[A\n",
      "  9%|▊         | 13/150 [00:05<00:53,  2.56it/s]\u001b[A\n",
      "  9%|▉         | 14/150 [00:05<00:53,  2.56it/s]\u001b[A\n",
      " 10%|█         | 15/150 [00:06<00:52,  2.56it/s]\u001b[A\n",
      " 11%|█         | 16/150 [00:06<00:52,  2.57it/s]\u001b[A\n",
      " 11%|█▏        | 17/150 [00:06<00:51,  2.57it/s]\u001b[A\n",
      " 12%|█▏        | 18/150 [00:07<00:51,  2.57it/s]\u001b[A\n",
      " 13%|█▎        | 19/150 [00:07<00:50,  2.57it/s]\u001b[A\n",
      " 13%|█▎        | 20/150 [00:07<00:50,  2.57it/s]\u001b[A\n",
      " 14%|█▍        | 21/150 [00:08<00:50,  2.57it/s]\u001b[A\n",
      " 15%|█▍        | 22/150 [00:08<00:49,  2.57it/s]\u001b[A\n",
      " 15%|█▌        | 23/150 [00:09<00:49,  2.57it/s]\u001b[A\n",
      " 16%|█▌        | 24/150 [00:09<00:48,  2.57it/s]\u001b[A\n",
      " 17%|█▋        | 25/150 [00:09<00:48,  2.57it/s]\u001b[A\n",
      " 17%|█▋        | 26/150 [00:10<00:48,  2.57it/s]\u001b[A\n",
      " 18%|█▊        | 27/150 [00:10<00:47,  2.57it/s]\u001b[A\n",
      " 19%|█▊        | 28/150 [00:11<00:47,  2.57it/s]\u001b[A\n",
      " 19%|█▉        | 29/150 [00:11<00:47,  2.57it/s]\u001b[A\n",
      " 20%|██        | 30/150 [00:11<00:46,  2.57it/s]\u001b[A\n",
      " 21%|██        | 31/150 [00:12<00:46,  2.57it/s]\u001b[A\n",
      " 21%|██▏       | 32/150 [00:12<00:45,  2.57it/s]\u001b[A\n",
      " 22%|██▏       | 33/150 [00:13<00:45,  2.57it/s]\u001b[A\n",
      " 23%|██▎       | 34/150 [00:13<00:45,  2.56it/s]\u001b[A\n",
      " 23%|██▎       | 35/150 [00:13<00:44,  2.57it/s]\u001b[A\n",
      " 24%|██▍       | 36/150 [00:14<00:44,  2.57it/s]\u001b[A\n",
      " 25%|██▍       | 37/150 [00:14<00:43,  2.57it/s]\u001b[A\n",
      " 25%|██▌       | 38/150 [00:14<00:43,  2.57it/s]\u001b[A\n",
      " 26%|██▌       | 39/150 [00:15<00:43,  2.57it/s]\u001b[A\n",
      " 27%|██▋       | 40/150 [00:15<00:42,  2.57it/s]\u001b[A\n",
      " 27%|██▋       | 41/150 [00:16<00:42,  2.56it/s]\u001b[A\n",
      " 28%|██▊       | 42/150 [00:16<00:42,  2.56it/s]\u001b[A\n",
      " 29%|██▊       | 43/150 [00:16<00:41,  2.56it/s]\u001b[A\n",
      " 29%|██▉       | 44/150 [00:17<00:41,  2.56it/s]\u001b[A\n",
      " 30%|███       | 45/150 [00:17<00:41,  2.56it/s]\u001b[A\n",
      " 31%|███       | 46/150 [00:18<00:40,  2.56it/s]\u001b[A\n",
      " 31%|███▏      | 47/150 [00:18<00:40,  2.56it/s]\u001b[A\n",
      " 32%|███▏      | 48/150 [00:18<00:39,  2.56it/s]\u001b[A\n",
      " 33%|███▎      | 49/150 [00:19<00:39,  2.56it/s]\u001b[A\n",
      " 33%|███▎      | 50/150 [00:19<00:40,  2.46it/s]\u001b[A\n",
      " 34%|███▍      | 51/150 [00:20<00:41,  2.40it/s]\u001b[A\n",
      " 35%|███▍      | 52/150 [00:20<00:41,  2.36it/s]\u001b[A\n",
      " 35%|███▌      | 53/150 [00:21<00:41,  2.33it/s]\u001b[A\n",
      " 36%|███▌      | 54/150 [00:21<00:41,  2.31it/s]\u001b[A\n",
      " 37%|███▋      | 55/150 [00:21<00:41,  2.30it/s]\u001b[A\n",
      " 37%|███▋      | 56/150 [00:22<00:41,  2.28it/s]\u001b[A\n",
      " 38%|███▊      | 57/150 [00:22<00:40,  2.28it/s]\u001b[A\n",
      " 39%|███▊      | 58/150 [00:23<00:40,  2.27it/s]\u001b[A\n",
      " 39%|███▉      | 59/150 [00:23<00:40,  2.27it/s]\u001b[A\n",
      " 40%|████      | 60/150 [00:24<00:39,  2.27it/s]\u001b[A\n",
      " 41%|████      | 61/150 [00:24<00:39,  2.27it/s]\u001b[A\n",
      " 41%|████▏     | 62/150 [00:24<00:38,  2.27it/s]\u001b[A\n",
      " 42%|████▏     | 63/150 [00:25<00:38,  2.27it/s]\u001b[A\n",
      " 43%|████▎     | 64/150 [00:25<00:37,  2.27it/s]\u001b[A\n",
      " 43%|████▎     | 65/150 [00:26<00:37,  2.27it/s]\u001b[A\n",
      " 44%|████▍     | 66/150 [00:26<00:37,  2.26it/s]\u001b[A\n",
      " 45%|████▍     | 67/150 [00:27<00:36,  2.27it/s]\u001b[A\n",
      " 45%|████▌     | 68/150 [00:27<00:36,  2.26it/s]\u001b[A\n",
      " 46%|████▌     | 69/150 [00:28<00:35,  2.26it/s]\u001b[A\n",
      " 47%|████▋     | 70/150 [00:28<00:35,  2.26it/s]\u001b[A\n",
      " 47%|████▋     | 71/150 [00:28<00:35,  2.26it/s]\u001b[A\n",
      " 48%|████▊     | 72/150 [00:29<00:34,  2.26it/s]\u001b[A\n",
      " 49%|████▊     | 73/150 [00:29<00:34,  2.26it/s]\u001b[A\n",
      " 49%|████▉     | 74/150 [00:30<00:33,  2.26it/s]\u001b[A\n",
      " 50%|█████     | 75/150 [00:30<00:33,  2.26it/s]\u001b[A\n",
      " 51%|█████     | 76/150 [00:31<00:32,  2.27it/s]\u001b[A\n",
      " 51%|█████▏    | 77/150 [00:31<00:32,  2.25it/s]\u001b[A\n",
      " 52%|█████▏    | 78/150 [00:32<00:31,  2.26it/s]\u001b[A\n",
      " 53%|█████▎    | 79/150 [00:32<00:31,  2.26it/s]\u001b[A\n",
      " 53%|█████▎    | 80/150 [00:32<00:31,  2.26it/s]\u001b[A\n",
      " 54%|█████▍    | 81/150 [00:33<00:30,  2.26it/s]\u001b[A\n",
      " 55%|█████▍    | 82/150 [00:33<00:30,  2.26it/s]\u001b[A\n",
      " 55%|█████▌    | 83/150 [00:34<00:29,  2.26it/s]\u001b[A\n",
      " 56%|█████▌    | 84/150 [00:34<00:29,  2.26it/s]\u001b[A\n",
      " 57%|█████▋    | 85/150 [00:35<00:28,  2.27it/s]\u001b[A\n",
      " 57%|█████▋    | 86/150 [00:35<00:28,  2.26it/s]\u001b[A\n",
      " 58%|█████▊    | 87/150 [00:36<00:27,  2.26it/s]\u001b[A\n",
      " 59%|█████▊    | 88/150 [00:36<00:27,  2.26it/s]\u001b[A\n",
      " 59%|█████▉    | 89/150 [00:36<00:27,  2.26it/s]\u001b[A\n",
      " 60%|██████    | 90/150 [00:37<00:26,  2.26it/s]\u001b[A\n",
      " 61%|██████    | 91/150 [00:37<00:26,  2.26it/s]\u001b[A\n",
      " 61%|██████▏   | 92/150 [00:38<00:25,  2.26it/s]\u001b[A\n",
      " 62%|██████▏   | 93/150 [00:38<00:25,  2.26it/s]\u001b[A\n",
      " 63%|██████▎   | 94/150 [00:39<00:24,  2.25it/s]\u001b[A\n",
      " 63%|██████▎   | 95/150 [00:39<00:24,  2.25it/s]\u001b[A\n",
      " 64%|██████▍   | 96/150 [00:40<00:23,  2.26it/s]\u001b[A\n",
      " 65%|██████▍   | 97/150 [00:40<00:23,  2.26it/s]\u001b[A\n",
      " 65%|██████▌   | 98/150 [00:40<00:22,  2.26it/s]\u001b[A\n",
      " 66%|██████▌   | 99/150 [00:41<00:22,  2.27it/s]\u001b[A\n",
      " 67%|██████▋   | 100/150 [00:41<00:22,  2.26it/s]\u001b[A\n",
      " 67%|██████▋   | 101/150 [00:42<00:21,  2.26it/s]\u001b[A\n",
      " 68%|██████▊   | 102/150 [00:42<00:21,  2.26it/s]\u001b[A\n",
      " 69%|██████▊   | 103/150 [00:43<00:20,  2.26it/s]\u001b[A\n",
      " 69%|██████▉   | 104/150 [00:43<00:20,  2.26it/s]\u001b[A\n",
      " 70%|███████   | 105/150 [00:44<00:19,  2.26it/s]\u001b[A\n",
      " 71%|███████   | 106/150 [00:44<00:19,  2.26it/s]\u001b[A\n",
      " 71%|███████▏  | 107/150 [00:44<00:19,  2.26it/s]\u001b[A\n",
      " 72%|███████▏  | 108/150 [00:45<00:18,  2.26it/s]\u001b[A\n",
      " 73%|███████▎  | 109/150 [00:45<00:18,  2.27it/s]\u001b[A\n",
      " 73%|███████▎  | 110/150 [00:46<00:17,  2.26it/s]\u001b[A\n",
      " 74%|███████▍  | 111/150 [00:46<00:17,  2.26it/s]\u001b[A\n",
      " 75%|███████▍  | 112/150 [00:47<00:16,  2.26it/s]\u001b[A\n",
      " 75%|███████▌  | 113/150 [00:47<00:16,  2.26it/s]\u001b[A\n",
      " 76%|███████▌  | 114/150 [00:47<00:15,  2.27it/s]\u001b[A\n",
      " 77%|███████▋  | 115/150 [00:48<00:15,  2.26it/s]\u001b[A\n",
      " 77%|███████▋  | 116/150 [00:48<00:15,  2.26it/s]\u001b[A\n",
      " 78%|███████▊  | 117/150 [00:49<00:14,  2.25it/s]\u001b[A\n",
      " 79%|███████▊  | 118/150 [00:49<00:14,  2.25it/s]\u001b[A\n",
      " 79%|███████▉  | 119/150 [00:50<00:13,  2.25it/s]\u001b[A\n",
      " 80%|████████  | 120/150 [00:50<00:13,  2.26it/s]\u001b[A\n",
      " 81%|████████  | 121/150 [00:51<00:12,  2.26it/s]\u001b[A\n",
      " 81%|████████▏ | 122/150 [00:51<00:12,  2.26it/s]\u001b[A\n",
      " 82%|████████▏ | 123/150 [00:51<00:11,  2.26it/s]\u001b[A\n",
      " 83%|████████▎ | 124/150 [00:52<00:11,  2.26it/s]\u001b[A\n",
      " 83%|████████▎ | 125/150 [00:52<00:11,  2.25it/s]\u001b[A\n",
      " 84%|████████▍ | 126/150 [00:53<00:10,  2.26it/s]\u001b[A\n",
      " 85%|████████▍ | 127/150 [00:53<00:10,  2.25it/s]\u001b[A\n",
      " 85%|████████▌ | 128/150 [00:54<00:09,  2.25it/s]\u001b[A\n",
      " 86%|████████▌ | 129/150 [00:54<00:09,  2.25it/s]\u001b[A\n",
      " 87%|████████▋ | 130/150 [00:55<00:08,  2.25it/s]\u001b[A\n",
      " 87%|████████▋ | 131/150 [00:55<00:08,  2.25it/s]\u001b[A\n",
      " 88%|████████▊ | 132/150 [00:55<00:07,  2.26it/s]\u001b[A\n",
      " 89%|████████▊ | 133/150 [00:56<00:07,  2.26it/s]\u001b[A\n",
      " 89%|████████▉ | 134/150 [00:56<00:07,  2.26it/s]\u001b[A\n",
      " 90%|█████████ | 135/150 [00:57<00:06,  2.26it/s]\u001b[A\n",
      " 91%|█████████ | 136/150 [00:57<00:06,  2.25it/s]\u001b[A\n",
      " 91%|█████████▏| 137/150 [00:58<00:05,  2.25it/s]\u001b[A\n",
      " 92%|█████████▏| 138/150 [00:58<00:05,  2.17it/s]\u001b[A\n",
      " 93%|█████████▎| 139/150 [00:59<00:05,  2.20it/s]\u001b[A\n",
      " 93%|█████████▎| 140/150 [00:59<00:04,  2.22it/s]\u001b[A\n",
      " 94%|█████████▍| 141/150 [01:00<00:04,  2.22it/s]\u001b[A\n",
      " 95%|█████████▍| 142/150 [01:00<00:03,  2.23it/s]\u001b[A\n",
      " 95%|█████████▌| 143/150 [01:00<00:03,  2.23it/s]\u001b[A\n",
      " 96%|█████████▌| 144/150 [01:01<00:02,  2.24it/s]\u001b[A\n",
      " 97%|█████████▋| 145/150 [01:01<00:02,  2.25it/s]\u001b[A\n",
      " 97%|█████████▋| 146/150 [01:02<00:01,  2.26it/s]\u001b[A\n",
      " 98%|█████████▊| 147/150 [01:02<00:01,  2.26it/s]\u001b[A\n",
      " 99%|█████████▊| 148/150 [01:03<00:00,  2.26it/s]\u001b[A\n",
      " 99%|█████████▉| 149/150 [01:03<00:00,  2.25it/s]\u001b[A\n",
      "100%|██████████| 150/150 [01:04<00:00,  2.34it/s]\u001b[A\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "thresh=100\n",
    "_ = [find_closest_people(frames[i],thresh,'frames2') for i in tqdm(range(len(frames))) ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 13.5 s, sys: 1.17 s, total: 14.7 s\n",
      "Wall time: 4.38 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "frames=[]\n",
    "for file in os.listdir(\"frames2/\"):\n",
    "    if file.endswith(\".png\"):\n",
    "        frames.append(file)\n",
    "frames.sort()\n",
    "\n",
    "frame_array=[]\n",
    "for i in range(len(frames)):\n",
    "    \n",
    "    #reading each files\n",
    "    img = cv2.imread('frames2/'+frames[i])\n",
    "    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n",
    "\n",
    "    height, width, layers = img.shape\n",
    "    size = (width,height)\n",
    "    \n",
    "    #inserting the frames into an image array\n",
    "    frame_array.append(img)\n",
    "\n",
    "out = cv2.VideoWriter('sample_output2.mp4',cv2.VideoWriter_fourcc(*'DIVX'), FPS, size)\n",
    " \n",
    "for i in range(len(frame_array)):\n",
    "    # writing to a image array\n",
    "    out.write(frame_array[i])\n",
    "out.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}