predict2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict\n",
    "\n",
    "**调用训练好的模型，预测图片中害虫位置及类别**\n",
    "\n",
    "**测试数据2**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 一、测试数据准备"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs_dir = '/opt/data/private/projects/TDCUP2022/datasets/tdcup/images2'\n",
    "\n",
    "print(len(os.listdir(imgs_dir)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 二、加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from detectron2.engine import DefaultPredictor\n",
    "from detectron2.config import get_cfg\n",
    "from mpvit import add_mpvit_config\n",
    "\n",
    "exp_tag = 'base_thred40_iter8w_data_dealed'\n",
    "\n",
    "cfg = get_cfg()\n",
    "add_mpvit_config(cfg)\n",
    "cfg.merge_from_file('configs/maskrcnn/mask_rcnn_mpvit_base_ms_3x.yaml')\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.4\n",
    "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 28\n",
    "cfg.MODEL.WEIGHTS = 'output/mask_rcnn_mpvit_base_ms_3x/exp2/model_0079999.pth'\n",
    "cfg.freeze()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Category ids in annotations are not in [1, #categories]! We'll apply a mapping for you.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['6', '7', '8', '9', '10', '25', '41', '105', '110', '115', '148', '156', '222', '228', '235', '256', '280', '310', '387', '392', '394', '398', '401', '402', '430', '480', '485', '673']\n"
     ]
    }
   ],
   "source": [
    "from detectron2.data import MetadataCatalog, DatasetCatalog\n",
    "\n",
    "\n",
    "# 获取数据集元数据\n",
    "pest_metadata = MetadataCatalog.get('pest_train')  #  property: ['json_file', 'name', 'set', 'thing_classes', 'thing_dataset_id_to_contiguous_id']\n",
    "pest_train_dataset = DatasetCatalog.get('pest_train')  # list of dict\n",
    "classes_name = pest_metadata.thing_classes\n",
    "\n",
    "print(classes_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建必要的文件夹\n",
    "\n",
    "output_dir = os.path.join('results/test2', exp_tag)\n",
    "imgs_save_dir = os.path.join(output_dir, 'pred_imgs')\n",
    "\n",
    "os.makedirs(output_dir, exist_ok=True)  # output path\n",
    "os.makedirs(imgs_save_dir, exist_ok=True)  # output path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 二、预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = DefaultPredictor(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "from detectron2.utils.visualizer import Visualizer\n",
    "\n",
    "import torch\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "result_list = []\n",
    "cnt = 1\n",
    "\n",
    "\n",
    "for idx, name_ in enumerate(os.listdir(imgs_dir)):\n",
    "    path_ = os.path.join(imgs_dir, name_)\n",
    "    img_save_path = os.path.join(imgs_save_dir, name_)\n",
    "\n",
    "    img = plt.imread(path_)\n",
    "    \n",
    "    if sum(img.shape) <= 3:\n",
    "        result_list.append([cnt, name_, 0, 0, 0, 0, 0, 0, 0])\n",
    "        cnt += 1\n",
    "        continue    \n",
    "\n",
    "    output = predictor(img)\n",
    "    instances = output[\"instances\"]\n",
    "\n",
    "    pred_classes = instances.pred_classes\n",
    "    pred_boxes = instances.pred_boxes\n",
    "\n",
    "    # continue when detect nothing\n",
    "    if pred_classes.shape[0] == 0:\n",
    "        result_list.append([cnt, name_, 0, 0, 0, 0, 0, 0, 0])\n",
    "        cnt += 1\n",
    "        continue\n",
    "\n",
    "    # save predicted images\n",
    "    v = Visualizer(img, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]))\n",
    "    out = v.draw_instance_predictions(instances.to(\"cpu\"))\n",
    "    pred_img = out.get_image()\n",
    "    \n",
    "    plt.imsave(img_save_path, pred_img)\n",
    "\n",
    "    # save predicted results to csv file\n",
    "    pred_classes = torch.tensor([int(classes_name[i]) for i in pred_classes], device=pred_classes.device)\n",
    "    for class_, box_ in zip(pred_classes, pred_boxes):\n",
    "        tpx, tpy, brx, bry = box_\n",
    "        mx, my = (tpx + brx) / 2, (tpy + bry) / 2\n",
    "        temp = [cnt, name_, class_.item(), mx.item(), my.item(), tpx.item(), tpy.item(), brx.item(), bry.item()]\n",
    "        result_list.append(temp)\n",
    "        cnt += 1\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 三、保存结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save result2.csv\n",
    "\n",
    "save_path = os.path.join(output_dir, 'result2.csv')\n",
    "\n",
    "df = pd.DataFrame(np.array(result_list))\n",
    "\n",
    "header = ['序号','文件名','虫子编号','中心点x坐标','中心点y坐标','左上角x坐标','左上角y坐标','右下角x坐标','右下角y坐标']\n",
    "\n",
    "df.to_csv(save_path, header=header,\n",
    "                        index=None, encoding='gbk')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save result3.csv\n",
    "\n",
    "save3_path = os.path.join(output_dir, 'result3.csv')\n",
    "\n",
    "res_dict = {}  # key为图像名称，value为字典（key为类别编号，value为该类别目标数量）\n",
    "for item_ in result_list:\n",
    "    name_, cid = item_[1:3]\n",
    "\n",
    "    if name_ not in res_dict.keys():\n",
    "        res_dict[name_] = {}\n",
    "    \n",
    "    if cid not in res_dict[name_].keys():\n",
    "        res_dict[name_][cid] = 0\n",
    "\n",
    "    if cid != 0:\n",
    "        res_dict[name_][cid] += 1\n",
    "\n",
    "\n",
    "res3 = []\n",
    "for name_, v in res_dict.items():\n",
    "    for kk, vv in v.items():\n",
    "        res3.append([name_, kk, vv])\n",
    "\n",
    "idx = (np.arange(len(res3)) + 1).reshape(-1, 1)\n",
    "\n",
    "data_new = np.hstack((idx, np.array(res3)))\n",
    "\n",
    "df_new = pd.DataFrame(data_new)\n",
    "\n",
    "header = ['序号','文件名','虫子编号','数量']\n",
    "\n",
    "df_new.to_csv(save3_path, header=header,\n",
    "                        index=False, encoding='gbk')"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "68c142f98e648b97731bcd2199f5f976913ca23975ebb19d28799c99c09961f5"
  },
  "kernelspec": {
   "display_name": "Python 3.8.8 ('detectron2')",
   "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.8.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}