{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Train a YOLO model for pose estimation on a custom keypoints dataset\n",
    "\n",
    "This notebook trains a YOLO model for pose estimation on the AnimalPose dataset.\n",
    "\n",
    "The input Table required for running this notebook is created in [create-custom-keypoints-table.ipynb](../1-create-tables/keypoints/create-custom-keypoints-table.ipynb).\n",
    "\n",
    "![](../images/animalpose-run.png)\n",
    "\n",
    "<!-- Tags: [\"ultralytics\", \"yolo\", \"keypoints\", \"animal-pose\", \"training\"] -->"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Project setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "PROJECT_NAME = \"3LC Tutorials - 2D Keypoints\"\n",
    "DATASET_NAME = \"AnimalPose\"\n",
    "TABLE_NAME = \"initial\"\n",
    "NUM_WORKERS = 0\n",
    "DOWNLOAD_PATH = \"../../transient_data\"\n",
    "EPOCHS = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install 3lc-ultralytics\n",
    "%pip install git+https://github.com/3lc-ai/3lc-examples.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tlc\n",
    "from tlc_ultralytics import YOLO, Settings\n",
    "\n",
    "from tlc_tools.split import split_table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Load and split table\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_table = tlc.Table.from_names(TABLE_NAME, DATASET_NAME, PROJECT_NAME)\n",
    "\n",
    "\n",
    "def split_by(table_row):\n",
    "    \"\"\"Callable to get the label of the first keypoint instance\n",
    "\n",
    "    This allows us to do a stratified split by label, just like in the original SuperGradients notebook.\n",
    "    \"\"\"\n",
    "    return table_row[\"keypoints_2d\"][\"instances_additional_data\"][\"label\"][0]\n",
    "\n",
    "\n",
    "train_val_test = split_table(\n",
    "    initial_table,\n",
    "    splits={\"train\": 0.8, \"val_test\": 0.2},\n",
    "    split_strategy=\"stratified\",\n",
    "    split_by=split_by,\n",
    "    random_seed=42,\n",
    "    shuffle=False,\n",
    ")\n",
    "\n",
    "test_val = split_table(\n",
    "    train_val_test[\"val_test\"],\n",
    "    splits={\"val\": 0.5, \"test\": 0.5},\n",
    "    split_strategy=\"stratified\",\n",
    "    split_by=split_by,\n",
    "    shuffle=False,\n",
    "    random_seed=42,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_table = train_val_test[\"train\"]\n",
    "val_table = test_val[\"val\"]\n",
    "test_table = test_val[\"test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(initial_table)\n",
    "print(train_table)\n",
    "print(val_table)\n",
    "print(test_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = YOLO(\"yolo11n-pose.pt\")\n",
    "\n",
    "settings = Settings(\n",
    "    project_name=PROJECT_NAME,\n",
    "    run_name=\"train-yolon-animalpose\",\n",
    "    run_description=\"Training a YOLO model for pose estimation on the AnimalPose dataset\",\n",
    "    collect_loss=True,\n",
    "    image_embeddings_dim=2,\n",
    ")\n",
    "\n",
    "model.train(\n",
    "    tables={\"train\": train_table, \"val\": val_table},\n",
    "    epochs=EPOCHS,\n",
    "    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"
  },
  "test_marks": [
   "dependent"
  ]
 },
 "nbformat": 4,
 "nbformat_minor": 5
}