Social-Distance-Tool-with-Depth.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",
    "from tqdm import tqdm\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": 9,
   "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",
    "import glob\n",
    "\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "\n",
    "import monodepth2.networks as networks\n",
    "from monodepth2.utils import download_model_if_doesnt_exist\n",
    "from monodepth2.layers import disp_to_depth\n",
    "import matplotlib as mpl\n",
    "import matplotlib.cm as cm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Define key variables**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames_folder = 'frames'\n",
    "result_folder = 'results'\n",
    "depths_folder = 'results/depth'\n",
    "frame_count = 10 # Number of frames to consider in the video (use less for faster calculations) # None will take all frames\n",
    "video = \"sample.mp4\"\n",
    "img_ext = 'png'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Convert Video to PNG Frames**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘frames/’: File exists\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:00<00:00, 18.82it/s]\n"
     ]
    }
   ],
   "source": [
    "!rm -r $frames_folder/*\n",
    "!mkdir $frames_folder/\n",
    "\n",
    "#specify path to video\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",
    "with tqdm(total=frame_count) as pbar:\n",
    "    while(cap.isOpened()):\n",
    "\n",
    "      # Capture frame-by-frame\n",
    "      ret, frame = cap.read()\n",
    "      pbar.update(1)\n",
    "      if ret == True:\n",
    "\n",
    "        #save each frame to folder        \n",
    "        cv2.imwrite(frames_folder+'/{:04d}'.format(cnt)+'.png', frame)\n",
    "        cnt=cnt+1\n",
    "        if(cnt==frame_count) and frame_count != None:\n",
    "          break\n",
    "      # Break the loop\n",
    "      else: \n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Loading MonoDept2 pretrained model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available(): # and not args.no_cuda:\n",
    "    device = torch.device(\"cuda\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "model_name = \"mono_640x192\"\n",
    "download_model_if_doesnt_exist(model_name)\n",
    "encoder_path = os.path.join(\"monodepth2/models\", model_name, \"encoder.pth\")\n",
    "depth_decoder_path = os.path.join(\"monodepth2/models\", model_name, \"depth.pth\")\n",
    "\n",
    "\n",
    "# LOADING PRETRAINED MODEL\n",
    "encoder = networks.ResnetEncoder(18, False)\n",
    "loaded_dict_enc = torch.load(encoder_path, map_location=device)\n",
    "filtered_dict_enc = {k: v for k, v in loaded_dict_enc.items() if k in encoder.state_dict()}\n",
    "encoder.load_state_dict(filtered_dict_enc)\n",
    "\n",
    "# extract the height and width of image that this model was trained with\n",
    "feed_height = loaded_dict_enc['height']\n",
    "feed_width = loaded_dict_enc['width']\n",
    "\n",
    "\n",
    "encoder.to(device)\n",
    "encoder.eval();\n",
    "\n",
    "# LOADING PRETRAINED MODEL\n",
    "depth_decoder = networks.DepthDecoder(num_ch_enc=encoder.num_ch_enc, scales=range(4))\n",
    "loaded_dict = torch.load(depth_decoder_path, map_location=device)\n",
    "depth_decoder.load_state_dict(loaded_dict)\n",
    "\n",
    "depth_decoder.to(device)\n",
    "depth_decoder.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Performing depth estimation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def findDepth(image_path,output_directory,ext):\n",
    "    # FINDING INPUT IMAGES\n",
    "    if os.path.isfile(image_path):\n",
    "        # Only testing on a single image\n",
    "        paths = [image_path]\n",
    "        #output_directory = os.path.dirname(args.image_path)\n",
    "    elif os.path.isdir(image_path):\n",
    "        # Searching folder for images\n",
    "        paths = glob.glob(os.path.join(image_path, '*.{}'.format(ext)))\n",
    "        #output_directory = args.image_path\n",
    "    else:\n",
    "        raise Exception(\"Can not find args.image_path: {}\".format(image_path))\n",
    "    #print(\"-> Predicting on {:d} test images\".format(len(paths)))\n",
    "\n",
    "\n",
    "    # PREDICTING ON EACH IMAGE IN TURN\n",
    "    with torch.no_grad():\n",
    "        with tqdm(total=len(paths)) as pbar:\n",
    "            for idx, image_path in (enumerate(paths)):\n",
    "\n",
    "                if image_path.endswith(\"_disp.jpg\"):\n",
    "                    # don't try to predict disparity for a disparity image!\n",
    "                    continue\n",
    "\n",
    "                # Load image and preprocess\n",
    "                input_image = pil.open(image_path).convert('RGB')\n",
    "                original_width, original_height = input_image.size\n",
    "                input_image = input_image.resize((feed_width, feed_height), pil.LANCZOS)\n",
    "                input_image = transforms.ToTensor()(input_image).unsqueeze(0)\n",
    "\n",
    "                # PREDICTION\n",
    "                input_image = input_image.to(device)\n",
    "                features = encoder(input_image)\n",
    "                outputs = depth_decoder(features)\n",
    "\n",
    "                disp = outputs[(\"disp\", 0)]\n",
    "                disp_resized = torch.nn.functional.interpolate(\n",
    "                    disp, (original_height, original_width), mode=\"bilinear\", align_corners=False)\n",
    "\n",
    "                # Saving numpy file\n",
    "                output_name = os.path.splitext(os.path.basename(image_path))[0]\n",
    "                name_dest_npy = os.path.join(output_directory, \"{}_disp.npy\".format(output_name))\n",
    "                scaled_disp, _ = disp_to_depth(disp, 0.1, 100)\n",
    "                np.save(name_dest_npy, scaled_disp.cpu().numpy())\n",
    "                \n",
    "                #print(np.shape(disp))\n",
    "                #break\n",
    "                \n",
    "                # Saving colormapped depth image\n",
    "                disp_resized_np = disp_resized.squeeze().cpu().numpy()\n",
    "                \n",
    "                #print(disp_resized_np)\n",
    "                \n",
    "                vmax = np.percentile(disp_resized_np, 95)\n",
    "                normalizer = mpl.colors.Normalize(vmin=disp_resized_np.min(), vmax=vmax)\n",
    "                \n",
    "                mapper = cm.ScalarMappable(norm=normalizer, cmap='magma')\n",
    "                colormapped_im = (mapper.to_rgba(disp_resized_np)[:, :, :3] * 255).astype(np.uint8)\n",
    "                \n",
    "                #print(np.shape(colormapped_im))\n",
    "                #break\n",
    "                \n",
    "                im = pil.fromarray(colormapped_im)\n",
    "                \n",
    "                name_dest_im = os.path.join(output_directory, \"{}_disp.jpeg\".format(output_name))\n",
    "                im.save(name_dest_im)\n",
    "                pbar.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘results/depth’: File exists\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:01<00:00,  6.53it/s]\n"
     ]
    }
   ],
   "source": [
    "!rm -r $depths_folder/*\n",
    "!mkdir $depths_folder\n",
    "findDepth(frames_folder,depths_folder,img_ext)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Download a pretrained model from Detectron2 Model Zoo for Faster-RCNN**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a function which return the bottom center of every bbox\n",
    "def mid_point(img,img_depth,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-1,y_mid-1)\n",
    "  #print(\"mid=\",mid)\n",
    "  #print(\"img_depth\",img_depth.size)\n",
    "  z_mid = img_depth[mid]\n",
    "  mid3d = (x_mid,y_mid,z_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",
    "  return mid3d\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,img_depth,person,p1,p2):\n",
    "  mid1 = []\n",
    "  mid2 = []\n",
    "  for p in p1:\n",
    "    mid1.append(mid_point(img,img_depth,person,p))\n",
    "  for pp in p2:\n",
    "    mid2.append(mid_point(img,img_depth,person,pp))\n",
    "  for inx in range(len(mid1)):\n",
    "      #print(\"mid1\",mid1[inx][:2])\n",
    "      _ = cv2.line(img, mid1[inx][:2], mid2[inx][:2], (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,name_depth,thresh,savedir):\n",
    "\n",
    "  img = cv2.imread(name)\n",
    "  depth = np.load(name_depth)\n",
    "  original_width, original_height,_ = img.shape # (1920,1080) #input_image.size\n",
    "  disp_resized = torch.nn.functional.interpolate(\n",
    "                    torch.from_numpy(depth), (original_height, original_width), mode=\"bilinear\", align_corners=False)\n",
    "  img_depth = disp_resized.squeeze().cpu().numpy()\n",
    "\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,img_depth,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,img_depth,person,p1,p2)\n",
    "  cv2.imwrite(savedir+'/'+name,img)\n",
    "  #print(savedir+'/'+name)\n",
    "  return 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Fetch all the frames of the video sequence**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames=[]\n",
    "for file in os.listdir(frames_folder):\n",
    "    if file.endswith(\".png\"):\n",
    "        frames.append(frames_folder+\"/\"+file)\n",
    "frames.sort()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Fetch all the frame depths of the video sequence**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_depths=[]\n",
    "for file in os.listdir(depths_folder):\n",
    "    if file.endswith(\".npy\"):\n",
    "        frame_depths.append(depths_folder+\"/\"+file)\n",
    "frame_depths.sort()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Main loop to get results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:04<00:00,  2.23it/s]\n"
     ]
    }
   ],
   "source": [
    "#from tqdm import tqdm\n",
    "thresh=110\n",
    "with tqdm(total=len(frames)) as pbar:\n",
    "    for i in range(len(frames)):\n",
    "        find_closest_people(frames[i],frame_depths[i],thresh,result_folder)\n",
    "        pbar.update(1)\n",
    "    \n",
    "#_ = [find_closest_people(frames[i],thresh,'frames2') for i in tqdm(range(len(frames))) ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Main file with highlighed results**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:00<00:00, 32.51it/s]\n",
      "100%|██████████| 10/10 [00:00<00:00, 104.33it/s]\n"
     ]
    }
   ],
   "source": [
    "frames=[]\n",
    "for file in os.listdir(result_folder+\"/frames/\"):\n",
    "    if file.endswith(\".png\"):\n",
    "        frames.append(file)\n",
    "frames.sort()\n",
    "\n",
    "frame_array=[]\n",
    "with tqdm(total=len(frames)) as pbar:\n",
    "    for i in range(len(frames)):\n",
    "\n",
    "        #reading each files\n",
    "        img = cv2.imread(result_folder+'/frames/'+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",
    "        pbar.update(1)\n",
    "\n",
    "out = cv2.VideoWriter(result_folder+'/result.mp4',cv2.VideoWriter_fourcc(*'DIVX'), FPS, size)\n",
    " \n",
    "with tqdm(total=len(frames)) as pbar:\n",
    "    for i in range(len(frame_array)):\n",
    "        # writing to a image array\n",
    "        out.write(frame_array[i])\n",
    "        pbar.update(1)\n",
    "out.release()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Depth map video**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:00<00:00, 68.44it/s]\n",
      "100%|██████████| 10/10 [00:00<00:00, 95.80it/s]\n"
     ]
    }
   ],
   "source": [
    "frames_depth=[]\n",
    "for file in os.listdir(depths_folder):\n",
    "    if file.endswith(\".jpeg\"):\n",
    "        frames_depth.append(file)\n",
    "frames_depth.sort()\n",
    "\n",
    "frame_array=[]\n",
    "with tqdm(total=len(frames_depth)) as pbar:\n",
    "    for i in range(len(frames_depth)):\n",
    "\n",
    "        #reading each files\n",
    "        img = cv2.imread(depths_folder+'/'+frames_depth[i])\n",
    "        #img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n",
    "        \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",
    "        pbar.update(1)\n",
    "\n",
    "out = cv2.VideoWriter(result_folder+'/depth_result.mp4',cv2.VideoWriter_fourcc(*'DIVX'), FPS, size)\n",
    " \n",
    "with tqdm(total=len(frames_depth)) as pbar:\n",
    "    for i in range(len(frame_array)):\n",
    "        # writing to a image array\n",
    "        out.write(frame_array[i])\n",
    "        pbar.update(1)\n",
    "out.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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
}