{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a YOLO segmentation model with 3LC metrics collection\n",
    "\n",
    "This notebook shows how to train a YOLO segmentation model with 3LC metrics collection on a YOLO-compatible 3LC `Table`.\n",
    "\n",
    "![](../images/cell-segmentations-run.png)\n",
    "\n",
    "<!-- Tags: [\"instance-segmentation\", \"ultralytics\", \"yolo\", \"training\", \"medical\"] -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install \"3lc[pacmap]\"\n",
    "%pip install 3lc-ultralytics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tlc\n",
    "from tlc_ultralytics import YOLO, Settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Project setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "PROJECT_NAME = \"3LC Tutorials - Cell Segmentation\"\n",
    "DATASET_NAME = \"Sartorius Cell Segmentation\"\n",
    "\n",
    "# Modify YOLO training parameters when training on your own data\n",
    "MODEL_NAME = \"yolo11n-seg.pt\"\n",
    "EPOCHS = 10\n",
    "BATCH_SIZE = 4\n",
    "DOWNLOAD_PATH = \"../../transient_data\"\n",
    "NUM_WORKERS = 8  # Multiple workers in notebook environment is not supported on Windows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_table = tlc.Table.from_names(\"train\", DATASET_NAME, PROJECT_NAME)\n",
    "val_table = tlc.Table.from_names(\"val\", DATASET_NAME, PROJECT_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = YOLO(MODEL_NAME)\n",
    "\n",
    "settings = Settings(\n",
    "    project_name=PROJECT_NAME,\n",
    "    run_name=\"Train YOLO Instance Segmentation Model\",\n",
    "    conf_thres=0.2,\n",
    "    sampling_weights=True,\n",
    "    exclude_zero_weight_training=True,\n",
    "    exclude_zero_weight_collection=False,\n",
    "    image_embeddings_dim=2,\n",
    ")\n",
    "\n",
    "results = model.train(\n",
    "    task=\"segment\",\n",
    "    tables={\n",
    "        \"train\": train_table,\n",
    "        \"val\": val_table,\n",
    "    },\n",
    "    settings=settings,\n",
    "    batch=BATCH_SIZE,\n",
    "    epochs=EPOCHS,\n",
    "    workers=NUM_WORKERS,\n",
    "    project=DOWNLOAD_PATH,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Run created at {results.run_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.12.10"
  },
  "test_marks": [
   "slow"
  ]
 },
 "nbformat": 4,
 "nbformat_minor": 2
}