demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams['figure.figsize'] = (25.0, 15.0)\n",
    "\n",
    "from detection import SSDDetector, plot_boxes\n",
    "from mscoco import class_names as coco_class_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "output_dir = './samples'\n",
    "param_file = './output/ssd_coco_mobilenet1.0_512x512-deploy-0150.params'\n",
    "model = SSDDetector(param_file, input_size=512)\n",
    "img_dir = os.path.expanduser('~/.mxnet/datasets/coco/val2017/')\n",
    "img_names = [img_name for img_name in os.listdir(img_dir)]\n",
    "img_names = img_names[10: 20]\n",
    "img_pathes = [os.path.join(img_dir, img_name) for img_name in img_names]\n",
    "imgs = [cv2.imread(img_path) for img_path in img_pathes]\n",
    "imgs = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in imgs]\n",
    "detections = model.detect(imgs)\n",
    "\n",
    "all_dets = []\n",
    "for img_dets in detections:\n",
    "    try:\n",
    "        img_dets = [[xmin, ymin, xmax, ymax, angle, coco_class_names[cls+1], score] for (xmin, ymin, xmax, ymax, angle, cls, score) in  img_dets]\n",
    "        all_dets.append(img_dets)\n",
    "    except:\n",
    "        raise ValueError ('Failed to get coco label with class_id {}'.format(cls))\n",
    "\n",
    "for img_name, img, detections in zip(img_names, imgs, all_dets):\n",
    "    print (img_name)\n",
    "    img = plot_boxes(img, detections)\n",
    "    plt.imshow(img)\n",
    "    plt.show()\n",
    "    \n",
    "    # save_path = os.path.join(output_dir, img_name)\n",
    "    # cv2.imwrite(save_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))"
   ]
  },
  {
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}