{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 代码实现"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由于Faster R-CNN代码架构过于庞大，我们将直接调用相关接口进行效果展示。先导入必要的包，并使用在MS COCO数据集上预训练的ResNet50作为主干网络。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T06:28:26.010951Z",
     "iopub.status.busy": "2023-01-10T06:28:26.010181Z",
     "iopub.status.idle": "2023-01-10T06:28:27.260312Z",
     "shell.execute_reply": "2023-01-10T06:28:27.259065Z",
     "shell.execute_reply.started": "2023-01-10T06:28:26.010824Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import functools\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T06:28:29.692266Z",
     "iopub.status.busy": "2023-01-10T06:28:29.691749Z",
     "iopub.status.idle": "2023-01-10T06:29:48.004555Z",
     "shell.execute_reply": "2023-01-10T06:29:48.003203Z",
     "shell.execute_reply.started": "2023-01-10T06:28:29.692235Z"
    }
   },
   "outputs": [],
   "source": [
    "! pip install -U 'git+https://github.com/MS COCOdataset/MS COCOapi.git#subdirectory=PythonAPI'\n",
    "! git clone https://github.com/pytorch/vision.git\n",
    "! cd vision;cp references/detection/utils.py ../\n",
    "! cp references/detection/transforms.py ../\n",
    "! cp references/detection/MS COCO_eval.py ../\n",
    "! cp references/detection/engine.py ../\n",
    "! cp references/detection/MS COCO_utils.py ../"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T06:30:28.868123Z",
     "iopub.status.busy": "2023-01-10T06:30:28.867452Z",
     "iopub.status.idle": "2023-01-10T06:30:41.051748Z",
     "shell.execute_reply": "2023-01-10T06:30:41.050718Z",
     "shell.execute_reply.started": "2023-01-10T06:30:28.868088Z"
    }
   },
   "outputs": [],
   "source": [
    "# 加载模型，使用MS COCO数据集预训练\n",
    "model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained='MS COCO')\n",
    "num_classes = 21  \n",
    "in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model.to(device)\n",
    "\n",
    "# 初始化优化器与学习率\n",
    "params = [p for p in model.parameters() if p.requires_grad]\n",
    "optimizer = torch.optim.SGD(params, lr=0.001, weight_decay=0.0005)\n",
    "lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "WEIGHTS_FILE = \"../input/fasterrcnn/faster_rcnn_state.pth\"\n",
    "model.load_state_dict(torch.load(WEIGHTS_FILE))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接着，编写模型测试的代码。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T06:31:22.873928Z",
     "iopub.status.busy": "2023-01-10T06:31:22.873574Z",
     "iopub.status.idle": "2023-01-10T06:31:22.884847Z",
     "shell.execute_reply": "2023-01-10T06:31:22.883595Z",
     "shell.execute_reply.started": "2023-01-10T06:31:22.873899Z"
    }
   },
   "outputs": [],
   "source": [
    "# 对模型进行测试\n",
    "def obj_detector(img):\n",
    "    img = cv2.imread(img, cv2.IMREAD_COLOR)\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)\n",
    "\n",
    "    # 导入图像并对图像进行处理\n",
    "    img /= 255.0\n",
    "    img = torch.from_numpy(img)\n",
    "    img = img.unsqueeze(0)\n",
    "    img = img.permute(0,3,1,2)\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    # 设置阈值\n",
    "    detection_threshold = 0.70\n",
    "    \n",
    "    img = list(im.to(device) for im in img)\n",
    "    output = model(img)\n",
    "\n",
    "    for i , im in enumerate(img):\n",
    "        boxes = output[i]['boxes'].data.cpu().numpy()\n",
    "        scores = output[i]['scores'].data.cpu().numpy()\n",
    "        labels = output[i]['labels'].data.cpu().numpy()\n",
    "        \n",
    "        labels = labels[scores >= detection_threshold]\n",
    "        boxes = boxes[scores >= detection_threshold].astype(np.int32)\n",
    "        scores = scores[scores >= detection_threshold]\n",
    "\n",
    "        boxes[:, 2] = boxes[:, 2] - boxes[:, 0]\n",
    "        boxes[:, 3] = boxes[:, 3] - boxes[:, 1]\n",
    "    \n",
    "    sample = img[0].permute(1,2,0).cpu().numpy()\n",
    "    sample = np.array(sample)\n",
    "    \n",
    "    boxes = output[0]['boxes'].data.cpu().numpy()\n",
    "    name = output[0]['labels'].data.cpu().numpy()\n",
    "    scores = output[0]['scores'].data.cpu().numpy()\n",
    "    \n",
    "    boxes = boxes[scores >= detection_threshold].astype(np.int32)\n",
    "    names = name.tolist()\n",
    "    \n",
    "    return names, boxes, sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在ImageNet上测试模型的效果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 984856 characters]",
      "text/plain": [
       "<Figure size 1440x4320 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pred_path = \"../input/imagenet/imagenet/val/\"\n",
    "pred_files = [os.path.join(pred_path,f) for f in os.listdir(pred_path)]\n",
    "\n",
    "classes= {1:'aeroplane', 2:'bicycle', 3:'bird', 4:'boat', 5:'bottle',\n",
    "          6:'bus', 7:'car', 8:'cat', 9:'chair', 10:'cow',\n",
    "          11:'diningtable', 12:'dog', 13:'horse', 14:'motorbike',\n",
    "          15:'person', 16:'pottedplant', 17:'sheep', 18:'sofa', \n",
    "          19:'train',20:'tvmonitor'}\n",
    "\n",
    "plt.figure(figsize=(20, 60))\n",
    "image_list = [0,11,17,28]\n",
    "for i, images in enumerate(pred_files):\n",
    "    if i > 30:\n",
    "        break\n",
    "    if i not in image_list:\n",
    "        continue\n",
    "\n",
    "    plt.subplot(10,2,image_list.index(i)+1)\n",
    "    \n",
    "    names, boxes, sample = obj_detector(images)\n",
    "    \n",
    "    for i,box in enumerate(boxes):\n",
    "        # 绘制包围盒\n",
    "        cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (0, 220, 0), 2)\n",
    "        cv2.putText(sample, classes[names[i]], (box[0],box[1]-5), cv2.FONT_HERSHEY_COMPLEX,\n",
    "                    0.7, (220,0,0), 1, cv2.LINE_AA)  \n",
    "\n",
    "    plt.axis('off')\n",
    "    plt.imshow(sample)\n"
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
