main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KYGsVqxrGupO"
   },
   "source": [
    "## Common Requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "r1TLc_zBDo1I",
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# download baseline code\n",
    "!git init\n",
    "!git pull https://github.com/epsilon-deltta/ssd_guillotine.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IFtobeRWDLPb"
   },
   "source": [
    "### Training On VOC2007 and VOC2012 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "0wgJkOwbw8dS",
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# download voc2007,2012 \n",
    "!wget http://pjreddie.com/media/files/VOCtrainval_06-Nov-2007.tar\n",
    "!wget http://pjreddie.com/media/files/VOCtest_06-Nov-2007.tar\n",
    "!wget http://pjreddie.com/media/files/VOCtrainval_11-May-2012.tar\n",
    "!wget http://pjreddie.com/media/files/VOC2012test.tar\n",
    "\n",
    "!tar -xvf VOCtrainval_06-Nov-2007.tar \n",
    "!tar -xvf VOCtest_06-Nov-2007.tar \n",
    "!tar -xvf VOCtrainval_11-May-2012.tar \n",
    "!tar -xvf VOC2012test.tar\n",
    "\n",
    "!rm  VOCtrainval_06-Nov-2007.tar\n",
    "!rm  VOCtest_06-Nov-2007.tar \n",
    "!rm  VOCtrainval_11-May-2012.tar \n",
    "!rm  VOC2012test.tar\n",
    "# integrated voc2007&voc2012 split data\n",
    "!python create_data_lists.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3MFofqjCT1pW"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim\n",
    "import torch.utils.data\n",
    "from model import SSD300, MultiBoxLoss\n",
    "from datasets import PascalVOCDataset\n",
    "from utils import *\n",
    "\n",
    "#### settings\n",
    "\n",
    "# Data parameters\n",
    "data_folder = './'  # folder with data files\n",
    "keep_difficult = True  # use objects considered difficult to detect?\n",
    "\n",
    "# Model parameters\n",
    "# Not too many here since the SSD300 has a very specific structure\n",
    "n_classes = len(label_map)  # number of different types of objects\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "# device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Learning parameters\n",
    "checkpoint = None  # path to model checkpoint, None if none\n",
    "batch_size = 8  # batch size\n",
    "iterations = 120000  # number of iterations to train\n",
    "workers = 2  # number of workers for loading data in the DataLoader\n",
    "print_freq = 200  # print training status every __ batches\n",
    "lr = 1e-3  # learning rate\n",
    "decay_lr_at = [80000, 100000]  # decay learning rate after these many iterations\n",
    "decay_lr_to = 0.1  # decay learning rate to this fraction of the existing learning rate\n",
    "momentum = 0.9  # momentum\n",
    "weight_decay = 5e-4  # weight decay\n",
    "grad_clip = None  # clip if gradients are exploding, which may happen at larger batch sizes (sometimes at 32) - you will recognize it by a sorting error in the MuliBox loss calculation\n",
    "\n",
    "cudnn.benchmark = True\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yNN3wJLrC8F-"
   },
   "outputs": [],
   "source": [
    "\n",
    "global start_epoch, label_map, epoch, checkpoint, decay_lr_at\n",
    "\n",
    "# Initialize model or load checkpoint\n",
    "if checkpoint is None:\n",
    "    start_epoch = 0\n",
    "    model = SSD300(n_classes=n_classes)\n",
    "    # Initialize the optimizer, with twice the default learning rate for biases, as in the original Caffe repo\n",
    "    biases = list()\n",
    "    not_biases = list()\n",
    "    for param_name, param in model.named_parameters():\n",
    "        if param.requires_grad:\n",
    "            if param_name.endswith('.bias'):\n",
    "                biases.append(param)\n",
    "            else:\n",
    "                not_biases.append(param)\n",
    "    optimizer = torch.optim.SGD(params=[{'params': biases, 'lr': 2 * lr}, {'params': not_biases}],\n",
    "                                lr=lr, momentum=momentum, weight_decay=weight_decay)\n",
    "\n",
    "else:\n",
    "    checkpoint = torch.load(checkpoint)\n",
    "    start_epoch = checkpoint['epoch'] + 1\n",
    "    print('\\nLoaded checkpoint from epoch %d.\\n' % start_epoch)\n",
    "    model = checkpoint['model']\n",
    "    optimizer = checkpoint['optimizer']\n",
    "\n",
    "\n",
    "# Custom dataloaders\n",
    "# output types of dataset\n",
    "# n : number of objects wrt. each Image. \n",
    "# (torch.Size([3, 300, 300]),\n",
    "#  torch.Size([n, 4]),  # Ground truth Box\n",
    "#  torch.Size([n]),     # Class\n",
    "#  torch.Size([n]))     # Difficulty\n",
    "train_dataset = PascalVOCDataset(data_folder,\n",
    "                                    split='train',\n",
    "                                    keep_difficult=keep_difficult)\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True,\n",
    "                                            collate_fn=train_dataset.collate_fn, num_workers=workers,\n",
    "                                            pin_memory=True)  # note that we're passing the collate function here\n",
    "\n",
    "# Move to default device\n",
    "model = model.to(device)\n",
    "criterion = MultiBoxLoss(priors_cxcy=model.priors_cxcy).to(device)\n",
    "\n",
    "# Calculate total number of epochs to train and the epochs to decay learning rate at (i.e. convert iterations to epochs)\n",
    "# To convert iterations to epochs, divide iterations by the number of iterations per epoch\n",
    "# The paper trains for 120,000 iterations with a batch size of 32, decays after 80,000 and 100,000 iterations\n",
    "epochs = iterations // (len(train_dataset) // 32)\n",
    "decay_lr_at = [it // (len(train_dataset) // 32) for it in decay_lr_at]\n",
    "\n",
    "from train import *\n",
    "# Epochs\n",
    "for epoch in range(start_epoch, epochs):\n",
    "\n",
    "    # Decay learning rate at particular epochs\n",
    "    if epoch in decay_lr_at:\n",
    "        adjust_learning_rate(optimizer, decay_lr_to)\n",
    "\n",
    "    # One epoch's training\n",
    "    train(train_loader=train_loader,\n",
    "            model=model,\n",
    "            criterion=criterion,\n",
    "            optimizer=optimizer,\n",
    "            epoch=epoch)\n",
    "\n",
    "    # Save checkpoint\n",
    "    save_checkpoint(epoch, model, optimizer)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fn3xKcbFGfyt"
   },
   "source": [
    "### Detect objects ! (trained on VOC1207)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2021-05-19 22:54:29--  https://docs.google.com/uc?export=download&confirm=z8Mz&id=1bvJfF6r_zYl2xZEpYXxgb7jLQHFZ01Qe\n",
      "Resolving docs.google.com (docs.google.com)... 172.217.175.78, 2404:6800:4004:81e::200e\n",
      "Connecting to docs.google.com (docs.google.com)|172.217.175.78|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Moved Temporarily\n",
      "Location: https://doc-00-8g-docs.googleusercontent.com/docs/securesc/7svoncoj1647lrtukuuc9qctka8vc7js/7us5anoj57569n7pge9gli01iomal2e9/1621432500000/17550492739458416425/07411496331195239303Z/1bvJfF6r_zYl2xZEpYXxgb7jLQHFZ01Qe?e=download [following]\n",
      "--2021-05-19 22:54:30--  https://doc-00-8g-docs.googleusercontent.com/docs/securesc/7svoncoj1647lrtukuuc9qctka8vc7js/7us5anoj57569n7pge9gli01iomal2e9/1621432500000/17550492739458416425/07411496331195239303Z/1bvJfF6r_zYl2xZEpYXxgb7jLQHFZ01Qe?e=download\n",
      "Resolving doc-00-8g-docs.googleusercontent.com (doc-00-8g-docs.googleusercontent.com)... 172.217.27.65, 2404:6800:4004:80a::2001\n",
      "Connecting to doc-00-8g-docs.googleusercontent.com (doc-00-8g-docs.googleusercontent.com)|172.217.27.65|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://docs.google.com/nonceSigner?nonce=2vlpdmue312ce&continue=https://doc-00-8g-docs.googleusercontent.com/docs/securesc/7svoncoj1647lrtukuuc9qctka8vc7js/7us5anoj57569n7pge9gli01iomal2e9/1621432500000/17550492739458416425/07411496331195239303Z/1bvJfF6r_zYl2xZEpYXxgb7jLQHFZ01Qe?e%3Ddownload&hash=68k2hoqehu5b8vaoq49dkhm4aekef26a [following]\n",
      "--2021-05-19 22:54:30--  https://docs.google.com/nonceSigner?nonce=2vlpdmue312ce&continue=https://doc-00-8g-docs.googleusercontent.com/docs/securesc/7svoncoj1647lrtukuuc9qctka8vc7js/7us5anoj57569n7pge9gli01iomal2e9/1621432500000/17550492739458416425/07411496331195239303Z/1bvJfF6r_zYl2xZEpYXxgb7jLQHFZ01Qe?e%3Ddownload&hash=68k2hoqehu5b8vaoq49dkhm4aekef26a\n",
      "Connecting to docs.google.com (docs.google.com)|172.217.175.78|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://doc-00-8g-docs.googleusercontent.com/docs/securesc/7svoncoj1647lrtukuuc9qctka8vc7js/7us5anoj57569n7pge9gli01iomal2e9/1621432500000/17550492739458416425/07411496331195239303Z/1bvJfF6r_zYl2xZEpYXxgb7jLQHFZ01Qe?e=download&nonce=2vlpdmue312ce&user=07411496331195239303Z&hash=2fmjml0r2brgfqrc5auute8b7fptu6bh [following]\n",
      "--2021-05-19 22:54:31--  https://doc-00-8g-docs.googleusercontent.com/docs/securesc/7svoncoj1647lrtukuuc9qctka8vc7js/7us5anoj57569n7pge9gli01iomal2e9/1621432500000/17550492739458416425/07411496331195239303Z/1bvJfF6r_zYl2xZEpYXxgb7jLQHFZ01Qe?e=download&nonce=2vlpdmue312ce&user=07411496331195239303Z&hash=2fmjml0r2brgfqrc5auute8b7fptu6bh\n",
      "Connecting to doc-00-8g-docs.googleusercontent.com (doc-00-8g-docs.googleusercontent.com)|172.217.27.65|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: unspecified [application/x-tar]\n",
      "Saving to: ‘voc_pretrined.pth.tar’\n",
      "\n",
      "voc_pretrined.pth.t     [            <=>     ] 100.46M  11.2MB/s    in 9.3s    \n",
      "\n",
      "2021-05-19 22:54:41 (10.8 MB/s) - ‘voc_pretrined.pth.tar’ saved [105336769]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import subprocess as sbp\n",
    "# download pretrained_ssd_model on voc (including weights and shape both) \n",
    "# i also wonder why author who made this pretrined model named filename *.tar , it's not .tar file\n",
    "\n",
    "FILENAME = 'voc_pretrined.pth.tar'\n",
    "FILEID = '1bvJfF6r_zYl2xZEpYXxgb7jLQHFZ01Qe'\n",
    "# model_down_cmd = 'wget --load-cookies ~/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies ~/cookies.txt --keep-session-cookies --no-check-certificate \"https://docs.google.com/uc?export=download&id={FILEID}\" -O- | sed -rn \"s/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p\")&id={FILEID}\" -O {FILENAME} && rm -rf ~/cookies.txt'\n",
    "if os.path.exists(FILENAME):\n",
    "    print(FILENAME,'is already here. dont have to download 😁')\n",
    "else :\n",
    "    !wget --load-cookies ~/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies ~/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id={FILEID}' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id={FILEID}\" -O {FILENAME} && rm -rf ~/cookies.txt\n",
    "#     print(str(sbp.check_output(model_down_cmd),'utf-8' )  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V72GIU4B7apQ"
   },
   "outputs": [],
   "source": [
    "# \n",
    "# Enjoy pretrained model .\n",
    "# \n",
    "import time\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim\n",
    "import torch.utils.data\n",
    "from model import SSD300, MultiBoxLoss\n",
    "from datasets import PascalVOCDataset\n",
    "from utils import *\n",
    "from torchvision import transforms\n",
    "from utils import *\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Load model checkpoint\n",
    "checkpoint = FILENAME\n",
    "checkpoint = torch.load(checkpoint,map_location = device)\n",
    "\n",
    "start_epoch = checkpoint['epoch'] + 1\n",
    "print('\\nLoaded checkpoint from epoch %d.\\n' % start_epoch)\n",
    "model = checkpoint['model']\n",
    "model = model.to(device)\n",
    "model.eval()\n",
    "\n",
    "# Transforms\n",
    "resize = transforms.Resize((300, 300))\n",
    "to_tensor = transforms.ToTensor()\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "\n",
    "\n",
    "def detect(model,original_image, min_score=0.2, max_overlap=0.5, top_k=200, suppress=None):\n",
    "    \"\"\"\n",
    "    Detect objects in an image with a trained SSD300, and visualize the results.\n",
    "\n",
    "    :param original_image: image, a PIL Image\n",
    "    :param min_score: minimum threshold for a detected box to be considered a match for a certain class\n",
    "    :param max_overlap: maximum overlap two boxes can have so that the one with the lower score is not suppressed via Non-Maximum Suppression (NMS)\n",
    "    :param top_k: if there are a lot of resulting detection across all classes, keep only the top 'k'\n",
    "    :param suppress: classes that you know for sure cannot be in the image or you do not want in the image, a list\n",
    "    :return: annotated image, a PIL Image\n",
    "    \"\"\"\n",
    "\n",
    "    # Transform\n",
    "    image = normalize(to_tensor(resize(original_image)))\n",
    "\n",
    "    # Move to default device\n",
    "    image = image.to(device)\n",
    "\n",
    "    # Forward prop.\n",
    "    predicted_locs, predicted_scores = model(image.unsqueeze(0))\n",
    "\n",
    "    # Detect objects in SSD output\n",
    "    det_boxes, det_labels, det_scores = model.detect_objects(predicted_locs, predicted_scores, min_score=min_score,\n",
    "                                                             max_overlap=max_overlap, top_k=top_k)\n",
    "\n",
    "    # Move detections to the CPU\n",
    "    det_boxes = det_boxes[0].to('cpu')\n",
    "\n",
    "    # Transform to original image dimensions\n",
    "    original_dims = torch.FloatTensor(\n",
    "        [original_image.width, original_image.height, original_image.width, original_image.height]).unsqueeze(0)\n",
    "    det_boxes = det_boxes * original_dims\n",
    "\n",
    "    # Decode class integer labels\n",
    "    det_labels = [rev_label_map[l] for l in det_labels[0].to('cpu').tolist()]\n",
    "\n",
    "    # If no objects found, the detected labels will be set to ['0.'], i.e. ['background'] in SSD300.detect_objects() in model.py\n",
    "    if det_labels == ['background']:\n",
    "        # Just return original image\n",
    "        return original_image\n",
    "\n",
    "    # Annotate\n",
    "    annotated_image = original_image\n",
    "    draw = ImageDraw.Draw(annotated_image)\n",
    "    font = ImageFont.load_default()\n",
    "    # font = ImageFont.truetype(\"arial.ttf\", 15)\n",
    "    # font = ImageFont.truetype(\"calibril.ttf\", 15)\n",
    "\n",
    "    # Suppress specific classes, if needed\n",
    "    for i in range(det_boxes.size(0)):\n",
    "        if suppress is not None:\n",
    "            if det_labels[i] in suppress:\n",
    "                continue\n",
    "\n",
    "        # Boxes\n",
    "        box_location = det_boxes[i].tolist()\n",
    "        draw.rectangle(xy=box_location, outline=label_color_map[det_labels[i]])\n",
    "        draw.rectangle(xy=[l + 1. for l in box_location], outline=label_color_map[\n",
    "            det_labels[i]])  # a second rectangle at an offset of 1 pixel to increase line thickness\n",
    "        # draw.rectangle(xy=[l + 2. for l in box_location], outline=label_color_map[\n",
    "        #     det_labels[i]])  # a third rectangle at an offset of 1 pixel to increase line thickness\n",
    "        # draw.rectangle(xy=[l + 3. for l in box_location], outline=label_color_map[\n",
    "        #     det_labels[i]])  # a fourth rectangle at an offset of 1 pixel to increase line thickness\n",
    "\n",
    "        # Text\n",
    "        text_size = font.getsize(det_labels[i].upper())\n",
    "        text_location = [box_location[0] + 2., box_location[1] - text_size[1]]\n",
    "        textbox_location = [box_location[0], box_location[1] - text_size[1], box_location[0] + text_size[0] + 4.,\n",
    "                            box_location[1]]\n",
    "        draw.rectangle(xy=textbox_location, fill=label_color_map[det_labels[i]])\n",
    "        draw.text(xy=text_location, text=det_labels[i].upper(), fill='white',\n",
    "                  font=font)\n",
    "    del draw\n",
    "\n",
    "    return annotated_image\n",
    "\n",
    "\n",
    "\n",
    "img_path = '/content/img/000022.jpg'\n",
    "original_image = Image.open(img_path, mode='r')\n",
    "original_image = original_image.convert('RGB')\n",
    "anno_img = detect(model,original_image, min_score=0.2, max_overlap=0.5, top_k=200)\n",
    "from IPython.display import display\n",
    "display(anno_img)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "suW0tPcYRDCW"
   },
   "outputs": [],
   "source": [
    "from  PIL import Image\n",
    "import requests as rq\n",
    "import io\n",
    "from PIL import  ImageDraw\n",
    "from IPython.display import display\n",
    "\n",
    "# why this has error\n",
    "imgurl ='https://t3.daumcdn.net/thumb/R720x0/?fname=http://t1.daumcdn.net/brunch/service/user/1e0Z/image/7ulVcf_Wh5VicfYz2sVq3v_z37Y.jpg' \n",
    "imgurl = 'https://s3.ap-northeast-2.amazonaws.com/img.kormedi.com/news/article/__icsFiles/artimage/2014/08/29/c_km601/515966_540.jpg'\n",
    "kokobyte = rq.get(imgurl).content\n",
    "\n",
    "koko = Image.open(io.BytesIO(kokobyte))\n",
    "anno_img = detect(model,koko, min_score=0.2, max_overlap=0.5, top_k=200)\n",
    "display(anno_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rIozsgtI7gzU"
   },
   "outputs": [],
   "source": [
    "\n",
    "\"\"\"\n",
    "Training.\n",
    "\"\"\"\n",
    "global start_epoch, label_map, epoch, checkpoint, decay_lr_at\n",
    "\n",
    "Initialize model or load checkpoint\n",
    "if checkpoint is None:\n",
    "    start_epoch = 0\n",
    "    model = SSD300(n_classes=n_classes)\n",
    "    # Initialize the optimizer, with twice the default learning rate for biases, as in the original Caffe repo\n",
    "    biases = list()\n",
    "    not_biases = list()\n",
    "    for param_name, param in model.named_parameters():\n",
    "        if param.requires_grad:\n",
    "            if param_name.endswith('.bias'):\n",
    "                biases.append(param)\n",
    "            else:\n",
    "                not_biases.append(param)\n",
    "    optimizer = torch.optim.SGD(params=[{'params': biases, 'lr': 2 * lr}, {'params': not_biases}],\n",
    "                                lr=lr, momentum=momentum, weight_decay=weight_decay)\n",
    "\n",
    "else:\n",
    "    checkpoint = torch.load(checkpoint)\n",
    "    start_epoch = checkpoint['epoch'] + 1\n",
    "    print('\\nLoaded checkpoint from epoch %d.\\n' % start_epoch)\n",
    "    model = checkpoint['model']\n",
    "    optimizer = checkpoint['optimizer']\n",
    "\n",
    "\n",
    "# Move to default device\n",
    "# model = model.to(device)\n",
    "criterion = MultiBoxLoss(priors_cxcy=model.priors_cxcy).to(device)\n",
    "\n",
    "# Custom dataloaders\n",
    "train_dataset = PascalVOCDataset(data_folder,\n",
    "                                    split='train',\n",
    "                                    keep_difficult=keep_difficult)\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True,\n",
    "                                            collate_fn=train_dataset.collate_fn, num_workers=workers,\n",
    "                                            pin_memory=True)  # note that we're passing the collate function here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U11WB1ZXwEDJ"
   },
   "source": [
    "## Training on Kaist Multispectral Pedestrain Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pONO8Ux_AAsc"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/epsilon-deltta/LkaistPdt.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2vqXk4fvEM0v"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim\n",
    "import torch.utils.data\n",
    "from model import SSD300, MultiBoxLoss\n",
    "from datasets import PascalVOCDataset\n",
    "from utils import *\n",
    "\n",
    "#### settings\n",
    "import os\n",
    "repo_root = os.path.abspath('./LkaistPdt')\n",
    "dt_root = os.path.abspath('./LkaistPdt/data')\n",
    "anno_dir = os.path.abspath(os.path.join(dt_root,'annotations'))\n",
    "img_dir  = os.path.abspath(os.path.join(dt_root,'images'))\n",
    "\n",
    "keep_difficult = False  # use objects considered difficult to detect?\n",
    "n_classes = 3  # number of different types of objects\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# Learning parameters\n",
    "checkpoint = None  # path to model checkpoint, None if none\n",
    "batch_size = 8  # batch size\n",
    "iterations = 120000  # number of iterations to train\n",
    "workers = 2  # number of workers for loading data in the DataLoader\n",
    "print_freq = 200  # print training status every __ batches\n",
    "lr = 1e-3  # learning rate\n",
    "decay_lr_at = [80000, 100000]  # decay learning rate after these many iterations\n",
    "decay_lr_to = 0.1  # decay learning rate to this fraction of the existing learning rate\n",
    "momentum = 0.9  # momentum\n",
    "weight_decay = 5e-4  # weight decay\n",
    "grad_clip = None  # clip if gradients are exploding, which may happen at larger batch sizes (sometimes at 32) - you will recognize it by a sorting error in the MuliBox loss calculation\n",
    "\n",
    "cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0LjMGouCCzAe"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "import torchvision\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision.transforms import  ToTensor\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "class KaistPdt(torch.utils.data.Dataset):\n",
    "\n",
    "    def __init__(self,root,split_file,split='train',transforms=None,target_transform=None):\n",
    "        self.root  = root\n",
    "        self.split = split\n",
    "\n",
    "        # rgb,thermal ,rgbTF,thermalTF\n",
    "        with open(split_file,'r') as f:\n",
    "            fnames = f.readlines()\n",
    "            fnames = [x.replace('\\n','') for x in fnames]\n",
    "        \n",
    "        self.paths = []\n",
    "        self.objects = []\n",
    "        for fname in fnames:\n",
    "\n",
    "            rgb_ther = {'rgb':True,'thermal':True}\n",
    "\n",
    "            \n",
    "            rgb_path =  os.path.join(root,'images',fname.replace('.txt','.jpg') )\n",
    "            thermal_path =  os.path.join(root,'images',fname.replace('.txt','.jpg').replace('visible','lwir') )\n",
    "            \n",
    "            rgb_ther['rgb']     = True if os.path.exists(rgb_path) else False\n",
    "            rgb_ther['thermal'] = True if os.path.exists(thermal_path) else False\n",
    "\n",
    "            if rgb_ther['rgb'] is False and rgb_ther['thermal'] is False : # when there is no any image.\n",
    "                continue\n",
    "            \n",
    "            # make annotation variable ,not images(cuz of size)  \n",
    "            \n",
    "            ann_path = os.path.join(self.root,'annotations',fname)\n",
    "\n",
    "            with open(ann_path,'r') as f:\n",
    "                lines = f.readlines()\n",
    "            if len(lines) > 1 :\n",
    "                obj ={}\n",
    "                obj['boxes']  = []\n",
    "                obj['labels'] = []\n",
    "\n",
    "                for line in lines[1:] :\n",
    "                    items = line.split(' ')\n",
    "\n",
    "                    if 'cyclist' in items[0]:   \n",
    "                        label = 1\n",
    "                    elif 'person' in items[0]:\n",
    "                        label = 2\n",
    "                    elif 'people' in items[0]:\n",
    "                        label = 3 \n",
    "                    xywh = list(int(x) for x in items[1:5] )\n",
    "\n",
    "                    obj['boxes'] .append(xywh)\n",
    "                    obj['labels'].append(label)\n",
    "                \n",
    "                obj['boxes'] = torch.tensor(obj['boxes'])\n",
    "                obj['labels'] = torch.tensor(obj['labels'])\n",
    "\n",
    "                self.objects.append(obj)\n",
    "            else :\n",
    "                continue\n",
    "\n",
    "            self.paths.append(fname)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.paths)\n",
    "\n",
    "    def __getitem__(self,idx):\n",
    "        \n",
    "        img          = self.get_img(idx)\n",
    "        boxes,labels = self.get_label(idx)\n",
    "        img , boxes,labels = transform(img,boxes,labels,split=self.split)\n",
    "        return img, boxes,labels\n",
    "\n",
    "    def get_img(self,idx):\n",
    "        fname = self.paths[idx]\n",
    "\n",
    "        rgb_ther = {'rgb':True,'thermal':True}\n",
    "        \n",
    "        rgb_path =  os.path.join(self.root,'images',fname.replace('.txt','.jpg') )\n",
    "        thermal_path =  os.path.join(self.root,'images',fname.replace('.txt','.jpg').replace('visible','lwir') )\n",
    "        \n",
    "        rgb_ther['rgb']     = True if os.path.exists(rgb_path) else False\n",
    "        rgb_ther['thermal'] = True if os.path.exists(thermal_path) else False\n",
    "\n",
    "        rgb     = ToTensor()(Image.open(rgb_path    ) ) if rgb_ther['rgb'] else torch.zeros([3, 512, 640])\n",
    "        thermal = ToTensor()(Image.open(thermal_path) ) if rgb_ther['rgb'] else torch.zeros([3, 512, 640])\n",
    "        \n",
    "        img = rgb+thermal\n",
    "        return img\n",
    "    def get_label(self,idx):\n",
    "        # class : cyclist, person ,people \n",
    "        # loc : numOfobjs x 4(class_xywh)\n",
    "        # cl  : numOfobjs\n",
    "        print(idx)\n",
    "        print(self.objects[idx])\n",
    "        boxes  = self.objects[idx]['boxes'] \n",
    "        labels = self.objects[idx]['labels']\n",
    "        return  boxes, labels \n",
    "        \n",
    "trdt  = KaistPdt(dt_root,os.path.join(repo_root,'split','train.txt') )\n",
    "valdt = KaistPdt(dt_root,os.path.join(repo_root,'split','val.txt'))\n",
    "tedt  = KaistPdt(dt_root,os.path.join(repo_root,'split','test.txt'))\n",
    "tr    = DataLoader(trdt,batch_size =16)\n",
    "val    = DataLoader(valdt,batch_size =16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vT52ot6h__hK"
   },
   "outputs": [],
   "source": [
    "\n",
    "global start_epoch, label_map, epoch, checkpoint, decay_lr_at\n",
    "\n",
    "# Initialize model or load checkpoint\n",
    "if checkpoint is None:\n",
    "    start_epoch = 0\n",
    "    model = SSD300(n_classes=n_classes)\n",
    "    # Initialize the optimizer, with twice the default learning rate for biases, as in the original Caffe repo\n",
    "    biases = list()\n",
    "    not_biases = list()\n",
    "    for param_name, param in model.named_parameters():\n",
    "        if param.requires_grad:\n",
    "            if param_name.endswith('.bias'):\n",
    "                biases.append(param)\n",
    "            else:\n",
    "                not_biases.append(param)\n",
    "    optimizer = torch.optim.SGD(params=[{'params': biases, 'lr': 2 * lr}, {'params': not_biases}],\n",
    "                                lr=lr, momentum=momentum, weight_decay=weight_decay)\n",
    "\n",
    "else:\n",
    "    checkpoint = torch.load(checkpoint)\n",
    "    start_epoch = checkpoint['epoch'] + 1\n",
    "    print('\\nLoaded checkpoint from epoch %d.\\n' % start_epoch)\n",
    "    model = checkpoint['model']\n",
    "    optimizer = checkpoint['optimizer']\n",
    "\n",
    "train_loader = tr\n",
    "# Move to default device\n",
    "model = model.to(device)\n",
    "criterion = MultiBoxLoss(priors_cxcy=model.priors_cxcy).to(device)\n",
    "\n",
    "epochs = iterations // (len(tr.dataset) // 32)\n",
    "decay_lr_at = [it // (len(tr.dataset) // 32) for it in decay_lr_at]\n",
    "\n",
    "from train import *\n",
    "# Epochs\n",
    "for epoch in range(start_epoch, epochs):\n",
    "\n",
    "    # Decay learning rate at particular epochs\n",
    "    if epoch in decay_lr_at:\n",
    "        adjust_learning_rate(optimizer, decay_lr_to)\n",
    "\n",
    "    # One epoch's training\n",
    "    train(train_loader=train_loader,\n",
    "            model=model,\n",
    "            criterion=criterion,\n",
    "            optimizer=optimizer,\n",
    "            epoch=epoch)\n",
    "\n",
    "    # Save checkpoint\n",
    "    save_checkpoint(epoch, model, optimizer)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "By7KOqveWcEi"
   },
   "source": [
    "## Testing on Pedestrian Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F9ofpOvc-9Mk"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "new_ssd.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}