Mask_R_CNN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Mask R-CNN",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Qb7U_vrB5VnA",
        "outputId": "40c6f528-d011-4900-a6f2-28a248fdcc86"
      },
      "source": [
        "cd drive/MyDrive/2021-Creative"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/content/drive/MyDrive/2021-Creative\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9rssgNES7kiL"
      },
      "source": [
        "import os\n",
        "import glob\n",
        "import json\n",
        "\n",
        "import cv2\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.patches as patches\n",
        "from pycocotools.coco import COCO\n",
        "\n",
        "import torch\n",
        "import torch.optim as optim\n",
        "from torch import nn, Tensor\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "import torchvision\n",
        "import torchvision.transforms.functional as TF\n",
        "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
        "from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3KzKD_l27rYW"
      },
      "source": [
        "classes = (\n",
        "    'top', 'blouse', 't-shirt', 'Knitted fabri', 'shirt', 'bra top', \n",
        "    'hood', 'blue jeans', 'pants', 'skirt', 'leggings', 'jogger pants', \n",
        "    'coat', 'jacket', 'jumper', 'padding jacket', 'best', 'kadigan', \n",
        "    'zip up', 'dress', 'jumpsuit')\n",
        "\n",
        "\n",
        "class FashionDataset(Dataset):\n",
        "    def __init__(self, path, transforms=None):\n",
        "        self.coco = COCO(path)\n",
        "        self.image_ids = list(self.coco.imgToAnns.keys())\n",
        "        self.transforms = transforms\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.image_ids)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        image_id = self.image_ids[idx]\n",
        "        file_name = self.coco.loadImgs(image_id)[0]['file_name']\n",
        "        file_name = f'/content/drive/MyDrive/2021-K_fashion/train_new_all/{file_name}'\n",
        "        image = Image.open(file_name).convert('RGB')\n",
        "\n",
        "        annot_ids = self.coco.getAnnIds(imgIds=image_id)\n",
        "        annots = [x for x in self.coco.loadAnns(annot_ids) if x['image_id'] == image_id]\n",
        "        \n",
        "        boxes = np.array([annot['bbox'] for annot in annots], dtype=np.float32)\n",
        "        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n",
        "        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n",
        "\n",
        "        labels = np.array([annot['category_id'] for annot in annots], dtype=np.int32)\n",
        "        masks = np.array([self.coco.annToMask(annot) for annot in annots], dtype=np.uint8)\n",
        "\n",
        "        area = np.array([annot['area'] for annot in annots], dtype=np.float32)\n",
        "        iscrowd = np.array([annot['iscrowd'] for annot in annots], dtype=np.uint8)\n",
        "\n",
        "        target = {\n",
        "            'boxes': boxes,\n",
        "            'masks': masks,\n",
        "            'labels': labels,\n",
        "            'area': area,\n",
        "            'iscrowd': iscrowd}\n",
        "        \n",
        "        if self.transforms is not None:\n",
        "            image, target = self.transforms(image, target)\n",
        "            \n",
        "        target['boxes'] = torch.as_tensor(target['boxes'], dtype=torch.float32)\n",
        "        target['masks'] = torch.as_tensor(target['masks'], dtype=torch.uint8)\n",
        "        target['labels'] = torch.as_tensor(target['labels'], dtype=torch.int64)\n",
        "        target['area'] = torch.as_tensor(target['area'], dtype=torch.float32)\n",
        "        target['iscrowd'] = torch.as_tensor(target['iscrowd'], dtype=torch.uint8)            \n",
        "\n",
        "        return image, target"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8H5tG70W7tGu"
      },
      "source": [
        "class Compose:\n",
        "    def __init__(self, transforms):\n",
        "        self.transforms = transforms\n",
        "\n",
        "    def __call__(self, image, target):\n",
        "        for transform in self.transforms:\n",
        "            image, target = transform(\n",
        "                image, target)\n",
        "\n",
        "        return image, target\n",
        "\n",
        "\n",
        "class Resize:\n",
        "    def __init__(self, size, interpolation=Image.BILINEAR):\n",
        "        self.size = size\n",
        "        self.interpolation = interpolation\n",
        "\n",
        "    def __call__(self, image, target):\n",
        "        w, h = image.size\n",
        "        image = image.resize(self.size)\n",
        "\n",
        "        _masks = target['masks'].copy()\n",
        "        masks = np.zeros((_masks.shape[0], self.size[0], self.size[1]))\n",
        "        \n",
        "        for i, v in enumerate(_masks):\n",
        "            v = Image.fromarray(v).resize(self.size, resample=Image.BILINEAR)\n",
        "            masks[i] = np.array(v, dtype=np.uint8)\n",
        "\n",
        "        target['masks'] = masks\n",
        "        target['boxes'][:, [0, 2]] *= self.size[0] / w\n",
        "        target['boxes'][:, [1, 3]] *= self.size[1] / h\n",
        "        \n",
        "        return image, target\n",
        "        \n",
        "\n",
        "class ToTensor:\n",
        "    def __call__(self, image, target):\n",
        "        image = TF.to_tensor(image)\n",
        "        \n",
        "        return image, target"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yWId8bll7vZu"
      },
      "source": [
        "model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)\n",
        "in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
        "model.roi_heads.box_predictor = FastRCNNPredictor(in_features, len(classes)+1)\n",
        "in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels\n",
        "hidden_layer = 256\n",
        "model.roi_heads.mask_predictor = MaskRCNNPredictor(\n",
        "    in_features_mask, hidden_layer, len(classes)+1)"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NOnG_QTy9l3f"
      },
      "source": [
        "def save_checkpoint(epoch, model, optimizer):\n",
        "    \"\"\"\n",
        "    Save model checkpoint.\n",
        "\n",
        "    :param epoch: epoch number\n",
        "    :param model: model\n",
        "    :param optimizer: optimizer\n",
        "    \"\"\"\n",
        "    state = {'epoch': epoch,\n",
        "             'model': model,\n",
        "             'optimizer': optimizer}\n",
        "    filename = 'checkpoint.tar'\n",
        "    torch.save(state, filename)"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dxrT2Adi7yin",
        "outputId": "21c24f36-6955-4df0-9290-421b31705bac"
      },
      "source": [
        "batch_size = 16\n",
        "lr = 1e-3\n",
        "max_size = 800\n",
        "num_workers = 2\n",
        "print_freq = 500\n",
        "num_epochs = 5\n",
        "checkpoint = \"checkpoint.tar\" # checkpoint path\n",
        "device = 'cuda:0'\n",
        "\n",
        "transforms_train = Compose([\n",
        "    Resize((max_size, max_size)),\n",
        "    ToTensor()])\n",
        "\n",
        "\n",
        "def collate_fn(batch):\n",
        "    return tuple(zip(*batch))\n",
        "\n",
        "# Load Checkpoint\n",
        "if checkpoint is None:\n",
        "    start_epoch = 0\n",
        "    params = [p for p in model.parameters() if p.requires_grad]\n",
        "    optimizer = optim.Adam(params, lr=lr, weight_decay=1e-5)\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",
        "# Data Loader\n",
        "dataset = FashionDataset('/content/drive/MyDrive/2021-K_fashion/train.json', transforms=transforms_train)\n",
        "train_loader = DataLoader(\n",
        "    dataset, batch_size=batch_size, shuffle=True, \n",
        "    num_workers=num_workers, collate_fn=collate_fn)\n",
        "\n",
        "model.to(device)\n",
        "\n",
        "\n",
        "def train_fn():\n",
        "    model.train()\n",
        "    for epoch in range(start_epoch, num_epochs):\n",
        "        for i, (images, targets) in enumerate(train_loader):\n",
        "            optimizer.zero_grad()\n",
        "            images = [image.to(device) for image in images]\n",
        "            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n",
        "\n",
        "            losses = model(images, targets)\n",
        "            loss = sum(loss for loss in losses.values())\n",
        "            \n",
        "            if i % print_freq == 0:\n",
        "              print(\n",
        "                  f\"{epoch}, {i}, C: {losses['loss_classifier'].item():.5f}, M: {losses['loss_mask'].item():.5f}, \"\\\n",
        "                  f\"B: {losses['loss_box_reg'].item():.5f}, O: {losses['loss_objectness'].item():.5f}, T: {loss.item():.5f}\")\n",
        "              \n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "        del images, targets, losses\n",
        "        save_checkpoint(epoch, model, optimizer)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Loaded checkpoint from epoch 1.\n",
            "\n",
            "loading annotations into memory...\n",
            "Done (t=1.88s)\n",
            "creating index...\n",
            "index created!\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fATRyZq08Qqd",
        "outputId": "b53ac294-c518-41d4-e569-de941de63df7"
      },
      "source": [
        "train_fn()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "1, 0, C: 0.63482, M: 0.66967, B: 0.01352, O: 0.02180, T: 1.48984\n",
            "1, 500, C: 0.43182, M: 0.65483, B: 0.11669, O: 0.04842, T: 1.25779\n",
            "1, 1000, C: 172.60710, M: 0.59062, B: 235.00777, O: 43.19550, T: 491.26004\n",
            "1, 1500, C: 1.09248, M: 0.65886, B: 0.30623, O: 0.25132, T: 2.56951\n",
            "1, 2000, C: 2.75686, M: 0.65495, B: 0.33983, O: 7.24888, T: 19.57147\n",
            "1, 2500, C: 0.56451, M: 0.65537, B: 0.12288, O: 0.04438, T: 1.39352\n",
            "1, 3000, C: 0.41165, M: 0.66035, B: 0.12372, O: 0.04163, T: 1.24574\n",
            "1, 3500, C: 0.32919, M: 0.66092, B: 0.11382, O: 0.01930, T: 1.12990\n",
            "1, 4000, C: 0.32494, M: 0.64419, B: 0.15721, O: 0.04611, T: 1.17864\n",
            "1, 4500, C: 0.23894, M: 0.66690, B: 0.10348, O: 0.04521, T: 1.06423\n",
            "1, 5000, C: 0.24160, M: 0.67656, B: 0.11298, O: 0.03652, T: 1.07305\n",
            "1, 5500, C: 0.26479, M: 0.68033, B: 0.13774, O: 0.03671, T: 1.12326\n",
            "2, 0, C: 0.23628, M: 0.63138, B: 0.12727, O: 0.03895, T: 1.04416\n",
            "2, 500, C: 0.20308, M: 0.65073, B: 0.09975, O: 7.51303, T: 8.47223\n",
            "2, 1000, C: 0.24078, M: 0.65304, B: 0.12874, O: 0.04527, T: 1.07563\n",
            "2, 1500, C: 0.29228, M: 0.64593, B: 0.16611, O: 0.02666, T: 1.13604\n",
            "2, 2000, C: 0.25013, M: 0.64805, B: 0.14483, O: 0.02542, T: 1.07642\n",
            "2, 2500, C: 10138.74707, M: 0.67926, B: 115.38312, O: 327696.62500, T: 1758627.75000\n",
            "2, 3000, C: 0.26303, M: 0.65162, B: 0.15143, O: 0.03818, T: 1.10974\n",
            "2, 3500, C: 0.23215, M: 0.67773, B: 0.12214, O: 0.03409, T: 1.06993\n",
            "2, 4000, C: 0.20848, M: 0.65982, B: 0.11571, O: 0.03833, T: 1.02835\n",
            "2, 4500, C: 0.23805, M: 0.66470, B: 0.13510, O: 0.04201, T: 1.08713\n",
            "2, 5000, C: 0.23756, M: 0.65840, B: 0.13156, O: 0.03609, T: 1.06901\n",
            "2, 5500, C: 0.21107, M: 0.66524, B: 0.10926, O: 0.03480, T: 1.02507\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZnEXa-NfupIc"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}