Day 60.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Detection using Single Shot Detection in Computer Vision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing the libraries\n",
    "import torch\n",
    "from torch.autograd import Variable\n",
    "import cv2\n",
    "from data import BaseTransform, VOC_CLASSES as labelmap\n",
    "from ssd import build_ssd\n",
    "import imageio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining a function that will do the detections\n",
    "def detect(frame, net, transform):\n",
    "    height, width = frame.shape[:2]\n",
    "    frame_t = transform(frame)[0]\n",
    "    x = torch.from_numpy(frame_t).permute(2, 0, 1)\n",
    "    x = Variable(x.unsqueeze(0))\n",
    "    y = net(x)\n",
    "    detections = y.data\n",
    "    scale = torch.Tensor([width, height, width, height])\n",
    "    # detections = [batch, number of classes, number of occurence, (score, x0, Y0, x1, y1)]\n",
    "    for i in range(detections.size(1)):\n",
    "        j = 0\n",
    "        while detections[0, i, j, 0] >= 0.6:\n",
    "            pt = (detections[0, i, j, 1:] * scale).numpy()\n",
    "            cv2.rectangle(frame, (int(pt[0]), int(pt[1])), (int(pt[2]), int(pt[3])), (255, 0, 0), 2)\n",
    "            cv2.putText(frame, labelmap[i - 1], (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 2, cv2.LINE_AA)\n",
    "            j += 1\n",
    "    return frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating the SSD neural network\n",
    "net = build_ssd('test')\n",
    "net.load_state_dict(torch.load('ssd300_mAP_77.43_v2.pth', map_location = lambda storage, loc: storage))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating the transformation\n",
    "transform = BaseTransform(net.size, (104/256.0, 117/256.0, 123/256.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\amira\\#100DaysOfMLCode\\Day 60\\ssd.py:100: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  self.softmax(conf.view(-1, self.num_classes)),  # conf preds\n",
      "WARNING:root:IMAGEIO FFMPEG_WRITER WARNING: input image is not divisible by macro_block_size=16, resizing from (1080, 1920) to (1088, 1920) to ensure video compatibility with most codecs and players. To prevent resizing, make your input image divisible by the macro_block_size or set the macro_block_size to None (risking incompatibility). You may also see a FFMPEG warning concerning speedloss due to data not being aligned.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n",
      "5\n",
      "6\n",
      "7\n",
      "8\n",
      "9\n",
      "10\n",
      "11\n",
      "12\n",
      "13\n",
      "14\n",
      "15\n",
      "16\n",
      "17\n",
      "18\n",
      "19\n",
      "20\n",
      "21\n",
      "22\n",
      "23\n",
      "24\n",
      "25\n",
      "26\n",
      "27\n",
      "28\n",
      "29\n",
      "30\n",
      "31\n",
      "32\n",
      "33\n",
      "34\n",
      "35\n",
      "36\n",
      "37\n",
      "38\n",
      "39\n",
      "40\n",
      "41\n",
      "42\n",
      "43\n",
      "44\n",
      "45\n",
      "46\n",
      "47\n",
      "48\n",
      "49\n",
      "50\n",
      "51\n",
      "52\n",
      "53\n",
      "54\n",
      "55\n",
      "56\n",
      "57\n",
      "58\n",
      "59\n",
      "60\n",
      "61\n",
      "62\n",
      "63\n",
      "64\n",
      "65\n",
      "66\n",
      "67\n"
     ]
    }
   ],
   "source": [
    "# Doing some Object Detection on a video\n",
    "reader = imageio.get_reader('funny_dog.mp4')\n",
    "fps = reader.get_meta_data()['fps']\n",
    "writer = imageio.get_writer('output.mp4', fps = fps)\n",
    "for i, frame in enumerate(reader):\n",
    "    frame = detect(frame, net.eval(), transform)\n",
    "    writer.append_data(frame)\n",
    "    print(i)\n",
    "writer.close()"
   ]
  }
 ],
 "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.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}