EfficientDet_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "GLb2lMFmNl8O"
   },
   "outputs": [],
   "source": [
    "# 라이브러리 및 모듈 import\n",
    "from pycocotools.coco import COCO\n",
    "from pycocotools.cocoeval import COCOeval\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import albumentations as A\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "from effdet import get_efficientdet_config, EfficientDet, DetBenchTrain\n",
    "from effdet.efficientdet import HeadNet\n",
    "import pandas as pd\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "eVbwnEcyNl8b"
   },
   "outputs": [],
   "source": [
    "# CustomDataset class 선언\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    '''\n",
    "      data_dir: data가 존재하는 폴더 경로\n",
    "      transforms: data transform (resize, crop, Totensor, etc,,,)\n",
    "    '''\n",
    "\n",
    "    def __init__(self, annotation, data_dir, transforms=None):\n",
    "        super().__init__()\n",
    "        self.data_dir = data_dir\n",
    "        # coco annotation 불러오기 (coco API)\n",
    "        self.coco = COCO(annotation)\n",
    "        self.predictions = {\n",
    "            \"images\": self.coco.dataset[\"images\"].copy(),\n",
    "            \"categories\": self.coco.dataset[\"categories\"].copy(),\n",
    "            \"annotations\": None\n",
    "        }\n",
    "        self.transforms = transforms\n",
    "\n",
    "    def __getitem__(self, index: int):\n",
    "        image_id = self.coco.getImgIds(imgIds=index)\n",
    "\n",
    "        image_info = self.coco.loadImgs(image_id)[0]\n",
    "        \n",
    "        image = cv2.imread(os.path.join(self.data_dir, image_info['file_name']))\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n",
    "        image /= 255.0\n",
    "\n",
    "        # 라벨 등 이미지 외 다른 정보 없기 때문에 train dataset과 달리 이미지만 전처리\n",
    "        \n",
    "        # transform\n",
    "        if self.transforms:\n",
    "            sample = self.transforms(image=image)\n",
    "\n",
    "        return sample['image'], image_id\n",
    "    \n",
    "    def __len__(self) -> int:\n",
    "        return len(self.coco.getImgIds())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "8XAO7td8Nl8d"
   },
   "outputs": [],
   "source": [
    "# Albumentation을 이용, augmentation 선언\n",
    "def get_train_transform():\n",
    "    return A.Compose([\n",
    "        A.Resize(512, 512),\n",
    "        A.Flip(p=0.5),\n",
    "        A.CLAHE(clip_limit=7.0,p=0.5),\n",
    "        A.RandomBrightness(limit=0.24, p=0.5),\n",
    "        ToTensorV2(p=1.0)\n",
    "    ])\n",
    "\n",
    "\n",
    "def get_valid_transform():\n",
    "    return A.Compose([\n",
    "        A.Resize(512, 512),\n",
    "        ToTensorV2(p=1.0)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "R_XLukpXNl8e"
   },
   "outputs": [],
   "source": [
    "from effdet import DetBenchPredict\n",
    "import gc\n",
    "\n",
    "# Effdet config를 통해 모델 불러오기 + ckpt load\n",
    "def load_net(checkpoint_path, device):\n",
    "    config = get_efficientdet_config('tf_efficientdet_d3')\n",
    "    config.num_classes = 10\n",
    "    config.image_size = (512,512)\n",
    "    \n",
    "    config.soft_nms = False\n",
    "    config.max_det_per_image = 25\n",
    "    \n",
    "    net = EfficientDet(config, pretrained_backbone=False)\n",
    "    net.class_net = HeadNet(config, num_outputs=config.num_classes)\n",
    "    \n",
    "    checkpoint = torch.load(checkpoint_path, map_location='cpu')\n",
    "\n",
    "    net = DetBenchPredict(net)\n",
    "    net.load_state_dict(checkpoint)\n",
    "    net.eval()\n",
    "\n",
    "    return net.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "zM1rE2zCNl8g"
   },
   "outputs": [],
   "source": [
    "# valid function\n",
    "def valid_fn(val_data_loader, model, device):\n",
    "    outputs = []\n",
    "    for images, image_ids in tqdm(val_data_loader):\n",
    "        # gpu 계산을 위해 image.to(device)       \n",
    "        images = torch.stack(images) # bs, ch, w, h \n",
    "        images = images.to(device).float()\n",
    "        output = model(images)\n",
    "        for out in output:\n",
    "            outputs.append({'boxes': out.detach().cpu().numpy()[:,:4], \n",
    "                            'scores': out.detach().cpu().numpy()[:,4], \n",
    "                            'labels': out.detach().cpu().numpy()[:,-1]})\n",
    "    return outputs\n",
    "\n",
    "def collate_fn(batch):\n",
    "    return tuple(zip(*batch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "eH4GRrXlNl8l"
   },
   "outputs": [],
   "source": [
    "def main():\n",
    "    annotation = './datasets/test.json'\n",
    "    data_dir = './datasets'\n",
    "    val_dataset = CustomDataset(annotation, data_dir, get_valid_transform())\n",
    "    epoch = 50\n",
    "    checkpoint_path = f'epoch_{epoch}.pth'\n",
    "    score_threshold = 0.1\n",
    "    val_data_loader = DataLoader(\n",
    "        val_dataset,\n",
    "        batch_size=4,\n",
    "        shuffle=False,\n",
    "        num_workers=4,\n",
    "        collate_fn=collate_fn\n",
    "    )\n",
    "    \n",
    "    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "    print(device)\n",
    "\n",
    "    model = load_net(checkpoint_path, device)\n",
    "    \n",
    "    outputs = valid_fn(val_data_loader, model, device)\n",
    "    \n",
    "    prediction_strings = []\n",
    "    file_names = []\n",
    "    coco = COCO(annotation)\n",
    "    \n",
    "    for i, output in enumerate(outputs):\n",
    "        prediction_string = ''\n",
    "        image_info = coco.loadImgs(coco.getImgIds(imgIds=i))[0]\n",
    "        for box, score, label in zip(output['boxes'], output['scores'], output['labels']):\n",
    "            if score > score_threshold:\n",
    "                 prediction_string += str(int(label)-1) + ' ' + str(score) + ' ' + str(box[0] * 2) + ' ' + str(\n",
    "                        box[1] * 2) + ' ' + str(box[2] * 2) + ' ' + str(box[3] * 2) + ' ' # 요기 \n",
    "        prediction_strings.append(prediction_string)\n",
    "        file_names.append(image_info['file_name'])\n",
    "        \n",
    "    submission = pd.DataFrame()\n",
    "    submission['PredictionString'] = prediction_strings\n",
    "    submission['image_id'] = file_names\n",
    "    submission.to_csv(f'submission_{epoch}.csv', index=None)\n",
    "    print(submission.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "1TpoSuGCNl8n",
    "outputId": "3df656f0-24df-4aa9-e544-3955810ef4c1",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.01s)\n",
      "creating index...\n",
      "index created!\n",
      "cuda\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1218/1218 [03:16<00:00,  6.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.01s)\n",
      "creating index...\n",
      "index created!\n",
      "                                    PredictionString       image_id\n",
      "0  7 0.9023277 606.5805053710938 505.516662597656...  test/0000.jpg\n",
      "1  3 0.6157351 351.52728271484375 254.15031433105...  test/0001.jpg\n",
      "2  1 0.7221721 767.0086669921875 515.344665527343...  test/0002.jpg\n",
      "3  9 0.7282079 116.83941650390625 246.00964355468...  test/0003.jpg\n",
      "4  1 0.71640575 224.30740356445312 306.9364624023...  test/0004.jpg\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TckxFyr3Nl8r"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "EfficientDet_inference.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}