add-bb-embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1c102552",
   "metadata": {},
   "source": [
    "# Per Bounding Box Embeddings Example"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3060059a",
   "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/add-bb-embeddings.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/add-bb-embeddings.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": "0edc8bc2",
   "metadata": {},
   "source": [
    "This notebook demonstrates how to extract embeddings for bounding boxes in `tlc` Tables using a pretrained EfficientNet model. \n",
    "The generated embeddings are then reduced in dimensionality and stored as extra columns in a new output Table.\n",
    "\n",
    "Since the example uses a classification model, we can also extract class probabilities for each bounding box.\n",
    "The predicted labels are also stored as additional columns in the Table.\n",
    "\n",
    "<div style=\"background-color: #e7f3fe; padding: 10px; border-left: 6px solid #2196F3; margin-bottom: 15px; color: #333;\">\n",
    "    <strong>Info:</strong> This notebook demonstrates a technique for adding columns to an existing Table.\n",
    "    While this is a useful technique, there are some drawbacks to this approach:\n",
    "    <ul>\n",
    "        <li>The new Table will not be part of the lineage of the input Table.</li>\n",
    "        <li>The new Table will contain a literal copy of all data in the input Table</li>\n",
    "    </ul>\n",
    "    In a future release of 3LC, adding columns to an existing Table will be supported natively.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b278eb8",
   "metadata": {},
   "source": [
    "## Project Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5575fc60",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "PROJECT_NAME = \"Bounding Box Embeddings\"\n",
    "DATASET_NAME = \"Balloons\"\n",
    "INSTALL_DEPENDENCIES = False\n",
    "TRANSIENT_DATA_PATH = \"../transient_data\"\n",
    "TEST_DATA_PATH = \"./data\"\n",
    "TLC_PUBLIC_EXAMPLES_DEVELOPER_MODE = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bc4ad42",
   "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 3lc[umap]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "473c3e2a",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3006bd36",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from io import BytesIO\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "\n",
    "import tlc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62e0365a",
   "metadata": {},
   "source": [
    "## Set Up Input Table\n",
    "\n",
    "We will use a `TableFromCoco` to load the \"Balloons\" dataset from a annotations file and a folder of images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81f5bd4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_url = tlc.Url.create_table_url(\n",
    "    project_name=PROJECT_NAME,\n",
    "    dataset_name=DATASET_NAME,\n",
    "    table_name=\"table_from_coco\",\n",
    ")\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 train split from COCO annotations\",\n",
    "    if_exists=\"overwrite\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ded6d8de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the schema of the bounding box column of the input table\n",
    "import json\n",
    "\n",
    "bb_schema = input_table.rows_schema.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 label(s): {json.dumps(label_map, indent=2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3986c01",
   "metadata": {},
   "source": [
    "## Initialize the Model\n",
    "\n",
    "Now we load the EfficientNet model. If a pretrained model is available locally, it will be loaded. Otherwise, we'll download a pretrained version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff377c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import timm\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Initialize a pretrained classifier model\n",
    "bb_classifier_path = TRANSIENT_DATA_PATH + \"/bb_classifier.pth\"\n",
    "if os.path.exists(bb_classifier_path):\n",
    "    model = timm.create_model(\"efficientnet_b0\", num_classes=2, checkpoint_path=bb_classifier_path).to(device)\n",
    "    print(\"Loaded pretrained model\")\n",
    "else:\n",
    "    print(\"Downloading pretrained model\")\n",
    "    model = timm.create_model(\"efficientnet_b0\", num_classes=len(label_map), pretrained=True).to(device)\n",
    "\n",
    "model = model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c82875b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The hidden layer we will use as embeddings\n",
    "hidden_layer = model.global_pool.flatten"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e365f06",
   "metadata": {},
   "source": [
    "## Collecting Bounding Box Embeddings\n",
    "\n",
    "In this section, we'll walk through the process of extracting embeddings for each bounding box present in our input images.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18a22c92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Image Preprocessing\n",
    "preprocess = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Initialize empty lists to store all embeddings and predicted labels\n",
    "all_embeddings: list[np.ndarray] = []\n",
    "all_labels: list[int] = []\n",
    "all_hidden_outputs: list[np.ndarray] = []\n",
    "\n",
    "# Register a hook to pick up the hidden layer output\n",
    "output_list: list[torch.Tensor] = []\n",
    "\n",
    "\n",
    "def hook_fn(module, input, output):\n",
    "    \"\"\"Store the output of the hooked layer.\"\"\"\n",
    "    output_list.append(output)\n",
    "\n",
    "\n",
    "hook_handle = hidden_layer.register_forward_hook(hook_fn)\n",
    "\n",
    "# Batched inference setup\n",
    "batch_size = 4\n",
    "mini_batch: list[torch.Tensor] = []\n",
    "batch_to_image_map: list[int] = []\n",
    "\n",
    "\n",
    "def run_inference_on_batch(mini_batch: list[torch.Tensor]) -> None:\n",
    "    mini_batch_tensor = torch.stack(mini_batch).to(device)\n",
    "    with torch.no_grad():\n",
    "        mini_batch_embeddings = model(mini_batch_tensor)\n",
    "\n",
    "    # Collect and clear the hook outputs\n",
    "    mini_batch_hidden = output_list.pop().cpu().numpy()\n",
    "    all_hidden_outputs.extend(mini_batch_hidden)\n",
    "\n",
    "    all_embeddings.extend(mini_batch_embeddings.cpu().numpy())\n",
    "    mini_batch_labels = torch.argmax(mini_batch_embeddings, dim=1)\n",
    "    all_labels.extend(mini_batch_labels.cpu().numpy())\n",
    "\n",
    "\n",
    "for row_idx, row in tqdm(enumerate(input_table.table_rows), total=len(input_table), desc=\"Running inference on table\"):\n",
    "    image_bbs = row[\"bbs\"][\"bb_list\"]\n",
    "    if len(image_bbs) == 0:\n",
    "        continue\n",
    "    image_filename = row[\"image\"]\n",
    "    image_bytes = tlc.Url(image_filename).read()\n",
    "    image = Image.open(BytesIO(image_bytes))\n",
    "    w, h = image.size\n",
    "\n",
    "    for bb in image_bbs:\n",
    "        bb_crop = tlc.BBCropInterface.crop(image, bb, bb_schema, h, w)\n",
    "        bb_crop_tensor = preprocess(bb_crop)\n",
    "\n",
    "        # Check if adding this bb_crop_tensor will overfill the mini_batch\n",
    "        if len(mini_batch) >= batch_size:\n",
    "            run_inference_on_batch(mini_batch)\n",
    "            mini_batch.clear()\n",
    "\n",
    "        mini_batch.append(bb_crop_tensor)\n",
    "        batch_to_image_map.append(row_idx)\n",
    "\n",
    "# Run inference on remaining items in mini_batch if it's not empty\n",
    "if len(mini_batch) > 0:\n",
    "    run_inference_on_batch(mini_batch)\n",
    "\n",
    "# Remove the hook\n",
    "hook_handle.remove()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d18b0b8c",
   "metadata": {},
   "source": [
    "## Dimensionality Reduction\n",
    "\n",
    "Once the embeddings are collected, the next step is to reduce their dimensionality for easier analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae6b6211",
   "metadata": {},
   "outputs": [],
   "source": [
    "import umap\n",
    "\n",
    "all_embeddings_np = np.vstack(all_hidden_outputs)\n",
    "print(f\"UMAP input shape: {all_embeddings_np.shape}\")\n",
    "\n",
    "# Fit UMAP\n",
    "reducer = umap.UMAP(n_components=3)\n",
    "embedding_3d = reducer.fit_transform(all_embeddings_np)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7bd0cda",
   "metadata": {},
   "source": [
    "## Create a new Table containing the embeddings as an extra column\n",
    "\n",
    "Finally, we combine the reduced embeddings and predicted labels with the input Table to write a new Table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4da749a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Repack embeddings and labels into groups per image\n",
    "grouped_embeddings: list[list[np.ndarray]] = [[] for _ in range(len(input_table))]\n",
    "grouped_labels: list[list[int]] = [[] for _ in range(len(input_table))]\n",
    "\n",
    "for img_idx, embed, label in zip(batch_to_image_map, embedding_3d, all_labels):\n",
    "    grouped_labels[img_idx].append(label)\n",
    "    grouped_embeddings[img_idx].append(embed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cfcf97d",
   "metadata": {},
   "source": [
    "### Setup the Schema of the output Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f73d0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a schema for the embeddings\n",
    "per_bb_embedding_schema = tlc.Schema(\n",
    "    value=tlc.Float32Value(number_role=tlc.NUMBER_ROLE_XYZ_COMPONENT),\n",
    "    size0=tlc.DimensionNumericValue(value_min=3, value_max=3),  # 3D embedding\n",
    "    size1=tlc.DimensionNumericValue(value_min=0, value_max=1000),  # Max 1000 bbs per image\n",
    "    sample_type=\"hidden\",  # Hide this column when iterating over the \"sample view\" of the table\n",
    "    writable=False,\n",
    ")\n",
    "\n",
    "# Create a schema with a label map for the labels\n",
    "label_value_map = {\n",
    "    **label_map,\n",
    "    **{len(label_map): tlc.MapElement(\"background\")},\n",
    "}\n",
    "\n",
    "label_schema = tlc.Schema(\n",
    "    value=tlc.Int32Value(value_map=label_value_map),\n",
    "    size0=tlc.DimensionNumericValue(value_min=0, value_max=1000),\n",
    "    sample_type=\"hidden\",  # Hide this column when iterating over the \"sample view\" of the table\n",
    "    writable=False,\n",
    ")\n",
    "\n",
    "schemas = {\n",
    "    \"per_bb_embeddings\": per_bb_embedding_schema,\n",
    "    \"per_bb_labels\": label_schema,\n",
    "}\n",
    "schemas.update(input_table.row_schema.values)  # Copy over the schemas from the input table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2b772c5",
   "metadata": {},
   "source": [
    "### Write the output Table\n",
    "\n",
    "We will use a `TableWriter` to write the output table as a `TableFromParquet`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49397f81",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "table_writer = tlc.TableWriter(\n",
    "    project_name=PROJECT_NAME,\n",
    "    dataset_name=DATASET_NAME,\n",
    "    description=\"Bounding box embeddings and labels\",\n",
    "    table_name=\"added_embeddings_and_labels\",\n",
    "    column_schemas=schemas,\n",
    "    if_exists=\"overwrite\",\n",
    ")\n",
    "\n",
    "# TableWriter accepts data as a dictionary of column names to lists\n",
    "data = defaultdict(list)\n",
    "\n",
    "# Copy over all rows from the input table\n",
    "for row in input_table.table_rows:\n",
    "    for column_name, column_value in row.items():\n",
    "        data[column_name].append(column_value)\n",
    "\n",
    "# Add the embeddings and labels\n",
    "data[\"per_bb_embeddings\"] = grouped_embeddings\n",
    "data[\"per_bb_labels\"] = grouped_labels\n",
    "\n",
    "table_writer.add_batch(data)\n",
    "new_table = table_writer.finalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec60061d",
   "metadata": {},
   "source": [
    "## Inspect the properties of the output Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "268af843",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(new_table))\n",
    "print(new_table.columns)\n",
    "print(new_table.url.to_relative(input_table.url))"
   ]
  }
 ],
 "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
}