test_on_val.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import cv2\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "import shutil\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mmdet.apis import init_detector, inference_detector, show_result_pyplot\n",
    "import mmcv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "shutil.rmtree('predictions', ignore_errors=True)\n",
    "os.mkdir('predictions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "BOX_COLOR = (0, 0, 255)\n",
    "TEXT_COLOR = (255, 255, 255)\n",
    "\n",
    "def visualize_bbox(img, bbox, \n",
    "#                    class_id, class_idx_to_name, \n",
    "                   score=None,\n",
    "                   color=BOX_COLOR, thickness=2):\n",
    "#     x_min, y_min, w, h = bbox\n",
    "    x_min, y_min, x_max, y_max = np.array(bbox).astype(np.int)\n",
    "#     x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(y_min + h)\n",
    "    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness)\n",
    "#     class_name = class_idx_to_name[class_id]\n",
    "    if score:\n",
    "        # pass bbox with small score\n",
    "        if score < 0.5:\n",
    "            return img\n",
    "        class_name = '%.3f' % score\n",
    "    else:\n",
    "        class_name = 'gt'\n",
    "    ((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1)    \n",
    "    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), color, -1)\n",
    "    cv2.putText(img, class_name, (x_min, y_min - int(0.3 * text_height)), cv2.FONT_HERSHEY_SIMPLEX, 0.35,TEXT_COLOR, lineType=cv2.LINE_AA)\n",
    "    return img\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file = 'mmdetection/configs/icartoonface/fr50_lite_dcn_gn_icf_ms49_1549_2x.py'\n",
    "# download the checkpoint from model zoo and put it in `checkpoints/`\n",
    "checkpoint_file = 'work_dirs/fr50_lite_dcn_gn_icf_ms49_1549_2x/epoch_21.pth'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build the model from a config file and a checkpoint file\n",
    "model = init_detector(config_file, checkpoint_file, device='cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load val set to test\n",
    "val_pkl  = '../data/icartoonface/dval.pkl'\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2500/2500 [07:17<00:00,  5.71it/s]\n"
     ]
    }
   ],
   "source": [
    "# visualize bbox gt and predictions on valdation data set\n",
    "data_val = mmcv.load(val_pkl)\n",
    "\n",
    "for data in tqdm(data_val):\n",
    "    file_name = data['filename']\n",
    "    img = cv2.imread('../data/' + file_name)\n",
    "    \n",
    "    \n",
    "    # draw gt bbox from annotations\n",
    "    for bbox in data['ann']['bboxes']:\n",
    "        img  = visualize_bbox(img, bbox, color=(0, 255, 0))\n",
    "        \n",
    "    # draw bbox predictions\n",
    "    result = inference_detector(model, img)\n",
    "    for bbox in result[0]:\n",
    "        img = visualize_bbox(img, bbox[:4], \n",
    "                            score=bbox[-1])\n",
    "        \n",
    "    \n",
    "        \n",
    "    # plot inline or save to disk\n",
    "#     plt.figure(figsize=(12, 12))\n",
    "#     plt.imshow(img[:,:,[2,1,0]])\n",
    "    cv2.imwrite('./predictions/' + file_name.split('/')[-1], img)"
   ]
  }
 ],
 "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.7.7"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}