main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import models\n",
    "from dataset import CocoDataset, collater, Resizer, AspectRatioBasedSampler, Augmenter, Normalizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = './data/'\n",
    "train_transforms = transforms.Compose([\n",
    "    Normalizer(), \n",
    "    Augmenter(), \n",
    "    Resizer()\n",
    "])\n",
    "val_transforms = transforms.Compose([\n",
    "    Normalizer(),\n",
    "    Resizer()\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(batch):\n",
    "    imgs = [s['img'] for s in batch]\n",
    "    annots = [s['annot'] for s in batch]\n",
    "\n",
    "    imgs = torch.stack([torch.from_numpy(img).permute(2, 0, 1) for img in imgs], 0)\n",
    "    annots = [torch.from_numpy(annot) for annot in annots]\n",
    "\n",
    "    return {'img': imgs, 'annot': annots}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.02s)\n",
      "creating index...\n",
      "index created!\n",
      "loading annotations into memory...\n",
      "Done (t=0.01s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "trainset = CocoDataset(root=path, set_name='train2017', transform=train_transforms, fraction=0.4)\n",
    "valset = CocoDataset(root=path, set_name='val2017', transform=val_transforms, fraction=0.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = AspectRatioBasedSampler(trainset, batch_size=2, drop_last=False)\n",
    "sampler_val = AspectRatioBasedSampler(valset, batch_size=1, drop_last=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainloader = DataLoader(trainset, num_workers=0, collate_fn=collater, batch_sampler=sampler, pin_memory=True)\n",
    "valloader = DataLoader(valset, num_workers=0, collate_fn=collater, batch_sampler=sampler_val, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "retinanet = models.resnet18(num_classes=trainset.num_classes())\n",
    "retinanet = torch.nn.DataParallel(retinanet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "retinanet.training = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "e:\\retinanet\\env\\Lib\\site-packages\\torch\\optim\\lr_scheduler.py:28: UserWarning: The verbose parameter is deprecated. Please use get_last_lr() to access the learning rate.\n",
      "  warnings.warn(\"The verbose parameter is deprecated. Please use get_last_lr() \"\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.AdamW(retinanet.parameters(), lr=1e-5)\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True)\n",
    "loss_hist = collections.deque(maxlen=500)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of training images : 100\n"
     ]
    }
   ],
   "source": [
    "retinanet.train()\n",
    "retinanet.module.freeze_bn()\n",
    "print(f'Number of training images : {len(trainloader)}');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0 | Iteration: 0 | Classification loss: 1.31181 | Regression loss: 1.02456 | Running loss: 2.33637\n",
      "Epoch: 0 | Iteration: 1 | Classification loss: 1.35089 | Regression loss: 1.03956 | Running loss: 2.36341\n",
      "Epoch: 0 | Iteration: 2 | Classification loss: 1.24268 | Regression loss: 1.01728 | Running loss: 2.32893\n",
      "Epoch: 0 | Iteration: 3 | Classification loss: 1.26288 | Regression loss: 1.04213 | Running loss: 2.32295\n",
      "Epoch: 0 | Iteration: 4 | Classification loss: 1.18546 | Regression loss: 1.07306 | Running loss: 2.31006\n",
      "Epoch: 0 | Iteration: 5 | Classification loss: 1.15298 | Regression loss: 1.03791 | Running loss: 2.29020\n",
      "Epoch: 0 | Iteration: 6 | Classification loss: 1.00447 | Regression loss: 1.01881 | Running loss: 2.25207\n",
      "Epoch: 0 | Iteration: 7 | Classification loss: 1.12124 | Regression loss: 1.01179 | Running loss: 2.23719\n",
      "Epoch: 0 | Iteration: 8 | Classification loss: 1.09549 | Regression loss: 1.07933 | Running loss: 2.23026\n",
      "Epoch: 0 | Iteration: 9 | Classification loss: 1.04520 | Regression loss: 1.03168 | Running loss: 2.21492\n",
      "Epoch: 0 | Iteration: 10 | Classification loss: 1.05253 | Regression loss: 1.04825 | Running loss: 2.20455\n",
      "Epoch: 0 | Iteration: 11 | Classification loss: 0.98811 | Regression loss: 1.07989 | Running loss: 2.19317\n",
      "Epoch: 0 | Iteration: 12 | Classification loss: 1.10504 | Regression loss: 1.05885 | Running loss: 2.19092\n",
      "Epoch: 0 | Iteration: 13 | Classification loss: 0.91502 | Regression loss: 1.08577 | Running loss: 2.17734\n",
      "Epoch: 0 | Iteration: 14 | Classification loss: 0.93348 | Regression loss: 1.12244 | Running loss: 2.16924\n",
      "Epoch: 0 | Iteration: 15 | Classification loss: 0.73400 | Regression loss: 1.09613 | Running loss: 2.14805\n",
      "Epoch: 0 | Iteration: 16 | Classification loss: 0.74221 | Regression loss: 1.04616 | Running loss: 2.12689\n",
      "Epoch: 0 | Iteration: 17 | Classification loss: 0.71197 | Regression loss: 1.02015 | Running loss: 2.10496\n",
      "Epoch: 0 | Iteration: 18 | Classification loss: 0.66357 | Regression loss: 0.96721 | Running loss: 2.08000\n",
      "Epoch: 0 | Iteration: 19 | Classification loss: 0.78075 | Regression loss: 1.03166 | Running loss: 2.06662\n",
      "Epoch: 0 | Iteration: 20 | Classification loss: 0.69886 | Regression loss: 1.03149 | Running loss: 2.05061\n",
      "Epoch: 0 | Iteration: 21 | Classification loss: 0.78790 | Regression loss: 1.05763 | Running loss: 2.04129\n",
      "Epoch: 0 | Iteration: 22 | Classification loss: 0.79959 | Regression loss: 0.89465 | Running loss: 2.02620\n",
      "Epoch: 0 | Iteration: 23 | Classification loss: 0.60428 | Regression loss: 0.93107 | Running loss: 2.00575\n",
      "Epoch: 0 | Iteration: 24 | Classification loss: 0.91377 | Regression loss: 0.98041 | Running loss: 2.00128\n",
      "Epoch: 0 | Iteration: 25 | Classification loss: 0.68499 | Regression loss: 1.23761 | Running loss: 1.99826\n",
      "Epoch: 0 | Iteration: 26 | Classification loss: 0.68076 | Regression loss: 0.95999 | Running loss: 1.98502\n",
      "Epoch: 0 | Iteration: 27 | Classification loss: 0.61351 | Regression loss: 0.97003 | Running loss: 1.97068\n",
      "Epoch: 0 | Iteration: 28 | Classification loss: 0.64249 | Regression loss: 0.91441 | Running loss: 1.95641\n",
      "Epoch: 0 | Iteration: 29 | Classification loss: 0.53242 | Regression loss: 0.96598 | Running loss: 1.94114\n",
      "Epoch: 0 | Iteration: 30 | Classification loss: 0.63050 | Regression loss: 0.98075 | Running loss: 1.93050\n",
      "Epoch: 0 | Iteration: 31 | Classification loss: 0.55048 | Regression loss: 0.99300 | Running loss: 1.91841\n",
      "Epoch: 0 | Iteration: 32 | Classification loss: 0.78658 | Regression loss: 0.90068 | Running loss: 1.91140\n",
      "Epoch: 0 | Iteration: 33 | Classification loss: 0.69931 | Regression loss: 0.96294 | Running loss: 1.90407\n",
      "Epoch: 0 | Iteration: 34 | Classification loss: 0.58904 | Regression loss: 1.03793 | Running loss: 1.89616\n",
      "Epoch: 0 | Iteration: 35 | Classification loss: 0.47665 | Regression loss: 0.95043 | Running loss: 1.88313\n",
      "Epoch: 0 | Iteration: 36 | Classification loss: 0.43604 | Regression loss: 0.90140 | Running loss: 1.86838\n",
      "Epoch: 0 | Iteration: 37 | Classification loss: 0.44239 | Regression loss: 0.97514 | Running loss: 1.85651\n",
      "Epoch: 0 | Iteration: 38 | Classification loss: 0.40885 | Regression loss: 0.86357 | Running loss: 1.84154\n",
      "Epoch: 0 | Iteration: 39 | Classification loss: 0.44253 | Regression loss: 1.05041 | Running loss: 1.83282\n",
      "Epoch: 0 | Iteration: 40 | Classification loss: 0.48961 | Regression loss: 0.93332 | Running loss: 1.82283\n",
      "Epoch: 0 | Iteration: 41 | Classification loss: 0.63232 | Regression loss: 1.20646 | Running loss: 1.82321\n",
      "Epoch: 0 | Iteration: 42 | Classification loss: 0.32394 | Regression loss: 0.98019 | Running loss: 1.81113\n",
      "Epoch: 0 | Iteration: 43 | Classification loss: 0.54986 | Regression loss: 0.84309 | Running loss: 1.80163\n",
      "Epoch: 0 | Iteration: 44 | Classification loss: 0.50500 | Regression loss: 0.81814 | Running loss: 1.79100\n",
      "Epoch: 0 | Iteration: 45 | Classification loss: 1.68928 | Regression loss: 0.63776 | Running loss: 1.80265\n",
      "Epoch: 0 | Iteration: 46 | Classification loss: 0.61357 | Regression loss: 1.07249 | Running loss: 1.80017\n",
      "Epoch: 0 | Iteration: 47 | Classification loss: 0.50356 | Regression loss: 1.04959 | Running loss: 1.79502\n",
      "Epoch: 0 | Iteration: 48 | Classification loss: 0.61359 | Regression loss: 0.91777 | Running loss: 1.78964\n",
      "Epoch: 0 | Iteration: 49 | Classification loss: 0.69719 | Regression loss: 1.02019 | Running loss: 1.78820\n",
      "Epoch: 0 | Iteration: 50 | Classification loss: 0.27350 | Regression loss: 0.89384 | Running loss: 1.77602\n",
      "Epoch: 0 | Iteration: 51 | Classification loss: 0.37281 | Regression loss: 0.92619 | Running loss: 1.76685\n",
      "Epoch: 0 | Iteration: 52 | Classification loss: 0.41511 | Regression loss: 0.88564 | Running loss: 1.75806\n",
      "Epoch: 0 | Iteration: 53 | Classification loss: 0.41806 | Regression loss: 0.91297 | Running loss: 1.75015\n",
      "Epoch: 0 | Iteration: 54 | Classification loss: 0.47248 | Regression loss: 1.08403 | Running loss: 1.74663\n",
      "Epoch: 0 | Iteration: 55 | Classification loss: 0.51919 | Regression loss: 1.13735 | Running loss: 1.74502\n",
      "Epoch: 0 | Iteration: 56 | Classification loss: 0.30601 | Regression loss: 0.88692 | Running loss: 1.73533\n",
      "Epoch: 0 | Iteration: 57 | Classification loss: 0.40078 | Regression loss: 1.09344 | Running loss: 1.73118\n",
      "Epoch: 0 | Iteration: 58 | Classification loss: 0.34520 | Regression loss: 0.90768 | Running loss: 1.72307\n",
      "Epoch: 0 | Iteration: 59 | Classification loss: 0.44452 | Regression loss: 0.93545 | Running loss: 1.71735\n",
      "Epoch: 0 | Iteration: 60 | Classification loss: 0.37527 | Regression loss: 0.96828 | Running loss: 1.71122\n",
      "Epoch: 0 | Iteration: 61 | Classification loss: 0.54420 | Regression loss: 0.87919 | Running loss: 1.70658\n",
      "Epoch: 0 | Iteration: 62 | Classification loss: 0.35807 | Regression loss: 0.92259 | Running loss: 1.69982\n",
      "Epoch: 0 | Iteration: 63 | Classification loss: 0.40582 | Regression loss: 1.04482 | Running loss: 1.69593\n",
      "Epoch: 0 | Iteration: 64 | Classification loss: 0.30925 | Regression loss: 0.91998 | Running loss: 1.68875\n",
      "Epoch: 0 | Iteration: 65 | Classification loss: 0.44833 | Regression loss: 1.15271 | Running loss: 1.68742\n",
      "Epoch: 0 | Iteration: 66 | Classification loss: 0.37696 | Regression loss: 0.95975 | Running loss: 1.68218\n",
      "Epoch: 0 | Iteration: 67 | Classification loss: 0.41959 | Regression loss: 0.94775 | Running loss: 1.67755\n",
      "Epoch: 0 | Iteration: 68 | Classification loss: 0.45550 | Regression loss: 1.10384 | Running loss: 1.67584\n",
      "Epoch: 0 | Iteration: 69 | Classification loss: 0.36299 | Regression loss: 0.87039 | Running loss: 1.66952\n",
      "Epoch: 0 | Iteration: 70 | Classification loss: 0.38871 | Regression loss: 0.99924 | Running loss: 1.66555\n",
      "Epoch: 0 | Iteration: 71 | Classification loss: 0.52477 | Regression loss: 0.95446 | Running loss: 1.66296\n",
      "Epoch: 0 | Iteration: 72 | Classification loss: 0.59574 | Regression loss: 1.05266 | Running loss: 1.66277\n",
      "Epoch: 0 | Iteration: 73 | Classification loss: 0.63863 | Regression loss: 1.11305 | Running loss: 1.66397\n",
      "Epoch: 0 | Iteration: 74 | Classification loss: 0.52372 | Regression loss: 1.00040 | Running loss: 1.66210\n",
      "Epoch: 0 | Iteration: 75 | Classification loss: 0.49148 | Regression loss: 0.98112 | Running loss: 1.65961\n",
      "Epoch: 0 | Iteration: 76 | Classification loss: 0.53968 | Regression loss: 0.93875 | Running loss: 1.65726\n",
      "Epoch: 0 | Iteration: 77 | Classification loss: 0.27192 | Regression loss: 1.00239 | Running loss: 1.65235\n",
      "Epoch: 0 | Iteration: 78 | Classification loss: 0.34065 | Regression loss: 1.04487 | Running loss: 1.64897\n",
      "Epoch: 0 | Iteration: 79 | Classification loss: 0.34594 | Regression loss: 0.89302 | Running loss: 1.64384\n",
      "Epoch: 0 | Iteration: 80 | Classification loss: 0.35129 | Regression loss: 0.95312 | Running loss: 1.63965\n",
      "Epoch: 0 | Iteration: 81 | Classification loss: 0.35036 | Regression loss: 1.00036 | Running loss: 1.63613\n",
      "Epoch: 0 | Iteration: 82 | Classification loss: 0.47577 | Regression loss: 1.01976 | Running loss: 1.63444\n",
      "Epoch: 0 | Iteration: 83 | Classification loss: 0.61331 | Regression loss: 1.00164 | Running loss: 1.63420\n",
      "Epoch: 0 | Iteration: 84 | Classification loss: 0.33909 | Regression loss: 0.97796 | Running loss: 1.63047\n",
      "Epoch: 0 | Iteration: 85 | Classification loss: 0.42687 | Regression loss: 1.03703 | Running loss: 1.62854\n",
      "Epoch: 0 | Iteration: 86 | Classification loss: 0.83681 | Regression loss: 1.05359 | Running loss: 1.63155\n",
      "Epoch: 0 | Iteration: 87 | Classification loss: 0.42568 | Regression loss: 0.89259 | Running loss: 1.62799\n",
      "Epoch: 0 | Iteration: 88 | Classification loss: 0.45373 | Regression loss: 0.97285 | Running loss: 1.62572\n",
      "Epoch: 0 | Iteration: 89 | Classification loss: 0.28380 | Regression loss: 0.96642 | Running loss: 1.62155\n",
      "Epoch: 0 | Iteration: 90 | Classification loss: 0.47034 | Regression loss: 0.84209 | Running loss: 1.61815\n",
      "Epoch: 0 | Iteration: 91 | Classification loss: 0.61292 | Regression loss: 0.89563 | Running loss: 1.61696\n",
      "Epoch: 0 | Iteration: 92 | Classification loss: 0.29647 | Regression loss: 0.85366 | Running loss: 1.61194\n",
      "Epoch: 0 | Iteration: 93 | Classification loss: 0.52398 | Regression loss: 1.01734 | Running loss: 1.61119\n",
      "Epoch: 0 | Iteration: 94 | Classification loss: 0.29621 | Regression loss: 0.85964 | Running loss: 1.60640\n",
      "Epoch: 0 | Iteration: 95 | Classification loss: 0.47102 | Regression loss: 0.93249 | Running loss: 1.60428\n",
      "Epoch: 0 | Iteration: 96 | Classification loss: 0.35218 | Regression loss: 0.93769 | Running loss: 1.60104\n",
      "Epoch: 0 | Iteration: 97 | Classification loss: 0.37437 | Regression loss: 0.98021 | Running loss: 1.59853\n",
      "Epoch: 0 | Iteration: 98 | Classification loss: 0.41773 | Regression loss: 0.96336 | Running loss: 1.59633\n",
      "Epoch: 0 | Iteration: 99 | Classification loss: 0.36879 | Regression loss: 0.93048 | Running loss: 1.59336\n"
     ]
    }
   ],
   "source": [
    "for epoch_num in range(num_epochs):\n",
    "\n",
    "        retinanet.train()\n",
    "        retinanet.module.freeze_bn()\n",
    "        epoch_loss = []\n",
    "\n",
    "        for iter_num, data in enumerate(trainloader):\n",
    "            try:\n",
    "                optimizer.zero_grad(set_to_none=True)\n",
    "                classification_loss, regression_loss = retinanet([data['img'].float(), data['annot']])    \n",
    "                classification_loss = classification_loss.mean()\n",
    "                regression_loss = regression_loss.mean()\n",
    "                loss = classification_loss + regression_loss\n",
    "                if bool(loss == 0):\n",
    "                    continue\n",
    "                loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(retinanet.parameters(), 0.1)\n",
    "                optimizer.step()\n",
    "                loss_hist.append(float(loss))\n",
    "                epoch_loss.append(float(loss))\n",
    "                print(\n",
    "                    'Epoch: {} | Iteration: {} | Classification loss: {:1.5f} | Regression loss: {:1.5f} | Running loss: {:1.5f}'.format(\n",
    "                        epoch_num, iter_num, float(classification_loss), float(regression_loss), np.mean(loss_hist)))\n",
    "                del classification_loss\n",
    "                del regression_loss\n",
    "            except Exception as e:\n",
    "                 print(e)\n",
    "                 continue\n",
    "            scheduler.step(np.mean(epoch_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(retinanet.module.state_dict(), 'retinanet_model_1.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycocotools.cocoeval import COCOeval\n",
    "import json\n",
    "\n",
    "def evaluate_coco(dataset, model, threshold=0.05):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        results = []\n",
    "        image_ids = []\n",
    "\n",
    "        for index in range(len(dataset)):\n",
    "            data = dataset[index]\n",
    "            scale = data['scale']\n",
    "\n",
    "            img = data['img'].unsqueeze(0).permute(0, 3, 1, 2).float()\n",
    "\n",
    "            if torch.cuda.is_available():\n",
    "                scores, labels, boxes = model(img.cuda())\n",
    "            else:\n",
    "                scores, labels, boxes = model(img)\n",
    "\n",
    "            scores = scores.cpu()\n",
    "            labels = labels.cpu()\n",
    "            boxes = boxes.cpu()\n",
    "\n",
    "            boxes /= scale\n",
    "\n",
    "            if boxes.shape[0] > 0:\n",
    "                boxes[:, 2] -= boxes[:, 0]\n",
    "                boxes[:, 3] -= boxes[:, 1]\n",
    "\n",
    "                for box_id in range(boxes.shape[0]):\n",
    "                    score = float(scores[box_id])\n",
    "                    label = int(labels[box_id])\n",
    "                    box = boxes[box_id, :]\n",
    "\n",
    "                    if score < threshold:\n",
    "                        break\n",
    "\n",
    "                    image_result = {\n",
    "                        'image_id': dataset.image_ids[index],\n",
    "                        'category_id': dataset.label_to_coco_label(label),\n",
    "                        'score': float(score),\n",
    "                        'bbox': box.tolist(),\n",
    "                    }\n",
    "                    results.append(image_result)\n",
    "\n",
    "            image_ids.append(dataset.image_ids[index])\n",
    "            print('{}/{}'.format(index, len(dataset)), end='\\r')\n",
    "\n",
    "        if not results:\n",
    "            return {\"error\": \"No detection results\"}\n",
    "\n",
    "        json.dump(results, open(f'{dataset.set_name}_bbox_results.json', 'w'), indent=4)\n",
    "\n",
    "        coco_true = dataset.coco\n",
    "        coco_pred = coco_true.loadRes(f'{dataset.set_name}_bbox_results.json')\n",
    "\n",
    "        coco_eval = COCOeval(coco_true, coco_pred, 'bbox')\n",
    "        coco_eval.params.imgIds = image_ids\n",
    "        coco_eval.evaluate()\n",
    "        coco_eval.accumulate()\n",
    "        coco_eval.summarize()\n",
    "\n",
    "        return coco_eval.stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}