train-bb-classifier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e50c4848",
   "metadata": {},
   "source": [
    "# Fine-tuning a Classifier Using Bounding Box Data from a 3LC Table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b917f8fc",
   "metadata": {},
   "source": [
    "<div style=\"display: inline-flex; align-items: center; gap: 10px;\">\n",
    "        <a href=\"https://colab.research.google.com/github/3lc-ai/notebook-examples/blob/main/train-bb-classifier.ipynb\"\n",
    "        target=\"_blank\"\n",
    "            style=\"background-color: transparent; text-decoration: none; display: inline-flex; align-items: center;\n",
    "            padding: 5px 10px; font-family: Arial, sans-serif;\"> <img\n",
    "            src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\" style=\"height: 30px;\n",
    "            vertical-align: middle;box-shadow: none;\"/>\n",
    "        </a> <a href=\"https://github.com/3lc-ai/notebook-examples/blob/main/train-bb-classifier.ipynb\"\n",
    "            style=\"text-decoration: none; display: inline-flex; align-items: center; background-color: #ffffff; border:\n",
    "            1px solid #d1d5da; border-radius: 8px; padding: 2px 10px; color: #333; font-family: Arial, sans-serif;\">\n",
    "            <svg aria-hidden=\"true\" focusable=\"false\" role=\"img\" class=\"octicon octicon-mark-github\" viewBox=\"0 0 16 16\"\n",
    "            width=\"20\" height=\"20\" fill=\"#333\"\n",
    "            style=\"display:inline-block;user-select:none;vertical-align:text-bottom;overflow:visible; margin-right:\n",
    "            8px;\">\n",
    "                <path d=\"M8 0c4.42 0 8 3.58 8 8a8.013 8.013 0 0 1-5.45 7.59c-.4.08-.55-.17-.55-.38 0-.27.01-1.13.01-2.2\n",
    "                0-.75-.25-1.23-.54-1.48 1.78-.2 3.65-.88 3.65-3.95 0-.88-.31-1.59-.82-2.15.08-.2.36-1.02-.08-2.12 0\n",
    "                0-.67-.22-2.2.82-.64-.18-1.32-.27-2-.27-.68 0-1.36.09-2 .27-1.53-1.03-2.2-.82-2.2-.82-.44 1.1-.16\n",
    "                1.92-.08 2.12-.51.56-.82 1.28-.82 2.15 0 3.06 1.86 3.75 3.64 3.95-.23.2-.44.55-.51\n",
    "                1.07-.46.21-1.61.55-2.33-.66-.15-.24-.6-.83-1.23-.82-.67.01-.27.38.01.53.34.19.73.9.82 1.13.16.45.68\n",
    "                1.31 2.69.94 0 .67.01 1.3.01 1.49 0 .21-.15.45-.55.38A7.995 7.995 0 0 1 0 8c0-4.42 3.58-8 8-8Z\"></path>\n",
    "            </svg> <span style=\"vertical-align: middle; color: #333;\">Open in GitHub</span>\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f66dd25e",
   "metadata": {},
   "source": [
    "This notebook demonstrates how to fine-tune a classifier on cropped bounding box regions. The goal is to aid in the evaluation of larger object detection networks. To achieve this, we implement a custom PyTorch DataLoader based on a single 3LC Table that contains image paths and associated bounding box information.\n",
    "\n",
    "We utilize a BBCropDataset, a specialized dataset class that ensures balanced sampling of bounding boxes for training. It tracks which bounding boxes have been sampled from each image to guarantee a well-distributed set of training examples.\n",
    "\n",
    "An optional feature allows the inclusion of a \"background\" class by sampling image regions devoid of any labeled bounding boxes. This is advantageous for enhancing the model's ability to discriminate between objects and background.\n",
    "\n",
    "The dataset supports on-the-fly data augmentation via PyTorch transforms and also provides an option to save cropped images, useful for both visualization and debugging.\n",
    "\n",
    "The model training employs the timm library for model selection. We pair the BBCropDataset with a RandomWeightedSampler to ensure each image is sampled in proportion to its number of bounding boxes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eb0d38b",
   "metadata": {},
   "source": [
    "## Project Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bc7d9c4",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "PROJECT_NAME = \"Bounding Box Classifier\"\n",
    "EPOCHS = 10\n",
    "TEST_DATA_PATH = \"./data\"\n",
    "TRANSIENT_DATA_PATH = \"../transient_data\"\n",
    "BATCH_SIZE = 32\n",
    "DATASET_NAME = \"Bounding Box Classification Dataset\"\n",
    "DEVICE = \"cuda:0\"\n",
    "TLC_PUBLIC_EXAMPLES_DEVELOPER_MODE = True\n",
    "INSTALL_DEPENDENCIES = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c97f513",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "if INSTALL_DEPENDENCIES:\n",
    "    %pip --quiet install torch --index-url https://download.pytorch.org/whl/cu118\n",
    "    %pip --quiet install torchvision --index-url https://download.pytorch.org/whl/cu118\n",
    "    %pip --quiet install timm\n",
    "    %pip --quiet install matplotlib\n",
    "    %pip --quiet install 3lc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc14b121",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03ce24ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from collections import defaultdict\n",
    "from io import BytesIO\n",
    "import random\n",
    "import os\n",
    "\n",
    "from PIL import Image\n",
    "import tqdm.notebook as tqdm\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "import tlc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a21b205a",
   "metadata": {},
   "source": [
    "## Set Up Input Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8cfd0dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_url = tlc.Url.create_table_url(project_name=PROJECT_NAME, dataset_name=DATASET_NAME, table_name=\"table_from_coco\")\n",
    "\n",
    "annotations_file = tlc.Url(TEST_DATA_PATH + \"/balloons/train/train-annotations.json\").to_absolute()\n",
    "images_dir = tlc.Url(TEST_DATA_PATH + \"/balloons/train\").to_absolute()\n",
    "\n",
    "input_table = tlc.Table.from_coco(\n",
    "    table_url=table_url,\n",
    "    annotations_file=annotations_file,\n",
    "    image_folder=images_dir,\n",
    "    description=\"Balloons training dataset from COCO annotations\",\n",
    "    if_exists=\"overwrite\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77f1bc49",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the columns of the input table\n",
    "print(input_table.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da92b766",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the schema of the bounding box column of the input table\n",
    "import json\n",
    "\n",
    "bb_schema = input_table.schema.values[\"rows\"].values[\"bbs\"].values[\"bb_list\"]\n",
    "label_map = input_table.get_value_map(\"bbs.bb_list.label\")\n",
    "print(f\"Input table uses {len(label_map)} unique labels: {json.dumps(label_map, indent=2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52212919",
   "metadata": {},
   "source": [
    "Notice our Table has a \"weight\" column which could be used for sampling, but since it is all 1's it is not very useful.\n",
    "\n",
    "We could however modify the weights in the Dashboard to control the sampling of images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b222645d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_weights_from_weight_column = [row[\"weight\"] for row in input_table.table_rows]\n",
    "print(sample_weights_from_weight_column)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac4dc203",
   "metadata": {},
   "source": [
    "We instead use the number of bounding boxes per image to control sampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4be24dba",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_weights_from_num_bbs = [len(row[\"bbs\"][\"bb_list\"]) for row in input_table.table_rows]\n",
    "print(sample_weights_from_num_bbs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46616a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import WeightedRandomSampler\n",
    "\n",
    "sampler = WeightedRandomSampler(sample_weights_from_num_bbs, num_samples=len(input_table))\n",
    "\n",
    "# Print the first 40 images to be drawn:\n",
    "print(\", \".join([str(next(iter(sampler))) for _ in range(40)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bedfe79",
   "metadata": {},
   "source": [
    "## Define Dataset\n",
    "\n",
    "Next, we define our dataset class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76c62f63",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BBCropDataset(Dataset):\n",
    "    def __init__(\n",
    "        self,\n",
    "        table: tlc.Table,\n",
    "        transform=None,\n",
    "        save_images_dir: str = \"\",\n",
    "        random_seed: int = 42,\n",
    "        is_train: bool = True,\n",
    "        background_freq: float = 0.5,\n",
    "    ):\n",
    "        self.table = table\n",
    "        self.bb_schema = table.schema.values[\"rows\"].values[\"bbs\"].values[\"bb_list\"]\n",
    "        self.transform = transform\n",
    "        self.save_images_dir = save_images_dir\n",
    "        self.used_bbs = defaultdict(set)\n",
    "        self.is_train = is_train\n",
    "        self.background_freq = background_freq\n",
    "        self.random_gen = random.Random(random_seed)\n",
    "\n",
    "        self.label_map = table.get_value_map(\"bbs.bb_list.label\")\n",
    "        self.background_label = len(self.label_map)\n",
    "\n",
    "        # Create mappings for contiguous IDs (in case label map is not contiguous)\n",
    "        self.id_to_contiguous = {\n",
    "            original_id: contiguous_id for contiguous_id, original_id in enumerate(self.label_map.keys())\n",
    "        }\n",
    "        self.contiguous_to_id = {\n",
    "            contiguous_id: original_id for original_id, contiguous_id in self.id_to_contiguous.items()\n",
    "        }\n",
    "        self.num_classes = len(self.label_map)\n",
    "\n",
    "        if background_freq > 0:\n",
    "            self.num_classes += 1  # Adding 1 for background class\n",
    "\n",
    "        if self.save_images_dir:\n",
    "            os.makedirs(self.save_images_dir, exist_ok=True)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.table)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        row = self.table.table_rows[idx]\n",
    "\n",
    "        image_filename = row[\"image\"]\n",
    "        image_bbs = row[\"bbs\"][\"bb_list\"]\n",
    "\n",
    "        if len(image_bbs) == 0:\n",
    "            raise ValueError(f\"Image {image_filename} has no bounding boxes. Use a sampler that excludes these images.\")\n",
    "\n",
    "        image_bytes = tlc.Url(image_filename).read()\n",
    "        image = Image.open(BytesIO(image_bytes))\n",
    "        w, h = image.size\n",
    "\n",
    "        available_bbs_idxs = list(set(range(len(image_bbs))) - self.used_bbs[idx])\n",
    "\n",
    "        if not available_bbs_idxs:\n",
    "            # print(f\"Re-using bbs from sample {idx}\")\n",
    "            self.used_bbs[idx] = set()\n",
    "            available_bbs_idxs = list(range(len(image_bbs)))\n",
    "\n",
    "        random_bb_idx = random.choice(available_bbs_idxs)\n",
    "\n",
    "        is_background = False\n",
    "        if self.random_gen.random() < self.background_freq and self.is_train:\n",
    "            is_background = True\n",
    "            gt_boxes = row[\"bbs\"][\"bb_list\"]\n",
    "            background_patch = self._generate_background(image, gt_boxes, w, h)\n",
    "            crop = background_patch\n",
    "            label = torch.tensor(self.background_label, dtype=torch.int64)\n",
    "        else:\n",
    "            random_bb = image_bbs[random_bb_idx]\n",
    "            self.used_bbs[idx].add(random_bb_idx)\n",
    "            crop = tlc.BBCropInterface.crop(image, random_bb, self.bb_schema, image_height=h, image_width=w)\n",
    "            label = torch.tensor(self.id_to_contiguous[random_bb[\"label\"]], dtype=torch.int64)\n",
    "\n",
    "        if self.save_images_dir:\n",
    "            crop.save(\n",
    "                os.path.join(self.save_images_dir, f\"{idx}_{random_bb_idx}{'_background' if is_background else ''}.jpg\")\n",
    "            )\n",
    "\n",
    "        if self.transform:\n",
    "            crop = self.transform(crop)\n",
    "\n",
    "        return crop, label\n",
    "\n",
    "    @staticmethod\n",
    "    def _intersects(box1: list[int], box2: list[int]) -> bool:\n",
    "        x1, y1, w1, h1 = box1\n",
    "        x2, y2, w2, h2 = box2\n",
    "\n",
    "        # Check for non-overlapping conditions\n",
    "        if x1 + w1 < x2 or x2 + w2 < x1 or y1 + h1 < y2 or y2 + h2 < y1:\n",
    "            return False\n",
    "        return True\n",
    "\n",
    "    def _generate_background(\n",
    "        self, image: Image.Image, gt_boxes: list, image_width: int, image_height: int\n",
    "    ) -> Image.Image:\n",
    "        \"\"\"Generate a background patch.\"\"\"\n",
    "        image_width, image_height = image.size\n",
    "        bb_factory = tlc.BoundingBox.from_schema(self.bb_schema)\n",
    "        gt_boxes_xywh = [\n",
    "            bb_factory([bb[\"x0\"], bb[\"y0\"], bb[\"x1\"], bb[\"y1\"]])\n",
    "            .to_top_left_xywh()\n",
    "            .denormalize(image_width, image_height)\n",
    "            for bb in gt_boxes\n",
    "        ]\n",
    "\n",
    "        # Loop until a valid background bounding box is generated\n",
    "        while True:\n",
    "            # Generate proposal box using normal distribution for x, h, w, y\n",
    "            x = max(\n",
    "                min(int(self.random_gen.normalvariate(mu=image_width // 2, sigma=image_width // 6)), image_width - 1), 0\n",
    "            )\n",
    "            y = max(\n",
    "                min(\n",
    "                    int(self.random_gen.normalvariate(mu=image_height // 2, sigma=image_height // 6)), image_height - 1\n",
    "                ),\n",
    "                0,\n",
    "            )\n",
    "            w = max(\n",
    "                min(int(self.random_gen.normalvariate(mu=image_width // 8, sigma=image_width // 16)), image_width - x),\n",
    "                1,\n",
    "            )\n",
    "            h = max(\n",
    "                min(\n",
    "                    int(self.random_gen.normalvariate(mu=image_height // 8, sigma=image_height // 16)), image_height - y\n",
    "                ),\n",
    "                1,\n",
    "            )\n",
    "\n",
    "            proposal_box = [x, y, w, h]\n",
    "\n",
    "            if not any(self._intersects(proposal_box, gt_box) for gt_box in gt_boxes_xywh):\n",
    "                break\n",
    "\n",
    "        # Crop the background patch from the image using the proposal_box\n",
    "        background_patch = image.crop((x, y, x + w, y + h))\n",
    "        return background_patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca33f67e",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_transforms = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(\n",
    "            mean=[0.485, 0.456, 0.406],\n",
    "            std=[0.229, 0.224, 0.225],\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "train_transforms = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.3),\n",
    "        transforms.RandomRotation(degrees=10),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(\n",
    "            mean=[0.485, 0.456, 0.406],\n",
    "            std=[0.229, 0.224, 0.225],\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "test_dataset = BBCropDataset(input_table, transform=train_transforms)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=2, sampler=sampler)\n",
    "\n",
    "# Read the first 3 batches\n",
    "num_test_batches = 3\n",
    "for batch_idx, batch in enumerate(test_dataloader):\n",
    "    images, labels = batch\n",
    "    print(f\"Received batch with images of shape {images.shape} and labels of shape {labels.shape}\")\n",
    "    if batch_idx == num_test_batches:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4efe92a",
   "metadata": {},
   "source": [
    "## Train a classifier using the crop dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d0398fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import timm\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "PRINT_FREQUENCY = 1\n",
    "\n",
    "# Initialize the model\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "train_dataset = BBCropDataset(input_table, transform=train_transforms, is_train=True)\n",
    "train_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, sampler=sampler)\n",
    "\n",
    "model = timm.create_model(\"efficientnet_b0\", pretrained=True, num_classes=train_dataset.num_classes).to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9516)\n",
    "\n",
    "best_accuracy = 0\n",
    "n_iter = 0\n",
    "iteration = []\n",
    "loss_history = []\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    total = 0\n",
    "    for i, (inputs, labels) in enumerate(tqdm.tqdm(train_dataloader)):\n",
    "        inputs, labels = inputs.to(device), labels.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "\n",
    "        total += labels.size(0)\n",
    "\n",
    "        n_iter += 1\n",
    "        if n_iter % PRINT_FREQUENCY == 0:\n",
    "            currentloss = 1.0 * running_loss / total\n",
    "            print(f\"Iteration: {n_iter}, Loss: {currentloss:.4f}\")\n",
    "            iteration.append(n_iter)\n",
    "            loss_history.append(currentloss)\n",
    "\n",
    "    if epoch >= 5:\n",
    "        scheduler.step()\n",
    "\n",
    "    print(f\"Epoch [{epoch+1}/{EPOCHS}], Loss: {running_loss / len(train_dataloader):.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a309047b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the model to a pth file:\n",
    "torch.save(model.state_dict(), TRANSIENT_DATA_PATH + \"/bb_classifier.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a0e2ebd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(iteration, loss_history)\n",
    "plt.title(\"Loss vs. Epoch\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}