{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a instance classifier on a 3LC Table\n",
    "\n",
    "In this tutorial, we will fine-tune a classifier using instances (segmentations or bounding boxes) from a 3LC `Table`.\n",
    "\n",
    "![](../../images/instance-embeddings.png)\n",
    "\n",
    "<!-- Tags: [\"object-detection\", \"classification\", \"embeddings\", \"metrics] -->\n",
    "\n",
    "We will load the COCO128 table from an earlier notebook and use it to create a\n",
    "`torch.utils.Dataset` of bounding box crops. These cropped images will be used to\n",
    "fine-tune a classifier. In a later tutorial, we will use this trained model to\n",
    "generate embeddings and predicted labels."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install 3lc[pacmap]\n",
    "%pip install git+https://github.com/3lc-ai/3lc-examples.git\n",
    "%pip install timm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tlc\n",
    "\n",
    "from tlc_tools.augment_bbs.finetune_on_crops import train_model\n",
    "from tlc_tools.common import infer_torch_device\n",
    "from tlc_tools.split import split_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Project setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "EPOCHS = 10\n",
    "DOWNLOAD_PATH = \"../../../transient_data\"\n",
    "MODEL_NAME = \"efficientnet_b0\"\n",
    "NUM_WORKERS = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_CHECKPOINT = DOWNLOAD_PATH + \"/instance_classifier.pth\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = infer_torch_device()\n",
    "print(f\"Using device: {DEVICE}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load input Table\n",
    "\n",
    "We will reuse the table created in the notebook [create-table-from-coco-detection.ipynb](../../1-create-tables/object%20detection/create-table-from-coco-detection.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_table = tlc.Table.from_names(\n",
    "    \"initial-segmentation\",\n",
    "    \"COCO128\",\n",
    "    \"3LC Tutorials - COCO128\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split the Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create splits for training and validation\n",
    "splits = split_table(input_table, {\"train\": 0.8, \"val\": 0.2}, if_exists=\"reuse\")\n",
    "\n",
    "train_table = splits[\"train\"]\n",
    "val_table = splits[\"val\"]\n",
    "\n",
    "print(f\"Using table {train_table} for training\")\n",
    "print(f\"Using table {val_table} for validation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, checkpoint_path = train_model(\n",
    "    train_table_url=train_table.url,\n",
    "    val_table_url=val_table.url,\n",
    "    model_checkpoint=MODEL_CHECKPOINT,\n",
    "    epochs=EPOCHS,\n",
    "    num_workers=NUM_WORKERS,\n",
    ")"
   ]
  }
 ],
 "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}