{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Object Detection with LTDETRv2\n",
    "\n",
    "This notebook demonstrates how to use LightlyTrain for object detection with our state-of-the-art LTDETRv2 model.\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/object_detection.ipynb)\n",
    "\n",
    "> **Important**: When running on Google Colab make sure to select a GPU runtime for faster processing. You can do this by going to `Runtime` > `Change runtime type` and selecting a GPU hardware accelerator."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "LightlyTrain can be installed directly via `pip`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install lightly-train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "> **Important**: LightlyTrain is officially supported on\n",
    "> - Linux: CPU or CUDA\n",
    "> - MacOS: CPU only\n",
    "> - Windows (experimental): CPU or CUDA\n",
    ">\n",
    "> We are planning to support MPS for MacOS.\n",
    ">\n",
    "> Check the [installation instructions](https://docs.lightly.ai/train/stable/installation.html) for more details on installation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Train an LTDETR model\n",
    "\n",
    "Training your own LTDETR model is straightforward with LightlyTrain.\n",
    "\n",
    "### Download a dataset\n",
    "\n",
    "First download a dataset. The dataset must be in YOLO format, see the [documentation](https://docs.lightly.ai/train/stable/object_detection.html#data) for more details. You can use [labelformat](https://github.com/lightly-ai/labelformat) to convert any dataset to the YOLO format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/lightly-ai/coco128_yolo/releases/download/v0.0.1/coco128_yolo.zip && unzip -q coco128_yolo.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "The dataset looks like this after the download completes:\n",
    "```\n",
    "coco128_yolo\n",
    "├── images\n",
    "│   ├── train2017\n",
    "│   │   ├── 000000000009.jpg\n",
    "│   │   ├── 000000000025.jpg\n",
    "│   │   ├── ...\n",
    "│   │   └── 000000000650.jpg\n",
    "│   └── val2017\n",
    "│       ├── 000000000139.jpg\n",
    "│       ├── 000000000285.jpg\n",
    "│       ├── ...\n",
    "│       └── 000000013201.jpg\n",
    "└── labels\n",
    "    ├── train2017\n",
    "    │   ├── 000000000009.txt\n",
    "    │   ├── 000000000025.txt\n",
    "    │   ├── ...\n",
    "    │   └── 000000000659.txt\n",
    "    └── val2017\n",
    "        ├── 000000000139.txt\n",
    "        ├── 000000000285.txt\n",
    "        ├── ...\n",
    "        └── 000000013201.txt\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "### Start training\n",
    "\n",
    "Next, start the training with the `train_object_detection` function. You only have to specify the output directory, model, and input data. LightlyTrain automatically sets the remaining training parameters and applies image augmentations. Of course you can always customize these settings if needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "lightly_train.train_object_detection(\n",
    "    out=\"out/my_experiment\",\n",
    "    model=\"ltdetrv2-s-coco\",\n",
    "    steps=100,  # Small number of steps for demonstration, default is 266_112.\n",
    "    batch_size=4,  # Small batch size for demonstration, default is 32.\n",
    "    data={\n",
    "        \"path\": \"coco128_yolo\",\n",
    "        \"train\": \"images/train2017\",\n",
    "        \"val\": \"images/val2017\",\n",
    "        \"names\": {\n",
    "            0: \"person\",\n",
    "            1: \"bicycle\",\n",
    "            2: \"car\",\n",
    "            3: \"motorcycle\",\n",
    "            4: \"airplane\",\n",
    "            5: \"bus\",\n",
    "            6: \"train\",\n",
    "            7: \"truck\",\n",
    "            8: \"boat\",\n",
    "            9: \"traffic light\",\n",
    "            10: \"fire hydrant\",\n",
    "            11: \"stop sign\",\n",
    "            12: \"parking meter\",\n",
    "            13: \"bench\",\n",
    "            14: \"bird\",\n",
    "            15: \"cat\",\n",
    "            16: \"dog\",\n",
    "            17: \"horse\",\n",
    "            18: \"sheep\",\n",
    "            19: \"cow\",\n",
    "            20: \"elephant\",\n",
    "            21: \"bear\",\n",
    "            22: \"zebra\",\n",
    "            23: \"giraffe\",\n",
    "            24: \"backpack\",\n",
    "            25: \"umbrella\",\n",
    "            26: \"handbag\",\n",
    "            27: \"tie\",\n",
    "            28: \"suitcase\",\n",
    "            29: \"frisbee\",\n",
    "            30: \"skis\",\n",
    "            31: \"snowboard\",\n",
    "            32: \"sports ball\",\n",
    "            33: \"kite\",\n",
    "            34: \"baseball bat\",\n",
    "            35: \"baseball glove\",\n",
    "            36: \"skateboard\",\n",
    "            37: \"surfboard\",\n",
    "            38: \"tennis racket\",\n",
    "            39: \"bottle\",\n",
    "            40: \"wine glass\",\n",
    "            41: \"cup\",\n",
    "            42: \"fork\",\n",
    "            43: \"knife\",\n",
    "            44: \"spoon\",\n",
    "            45: \"bowl\",\n",
    "            46: \"banana\",\n",
    "            47: \"apple\",\n",
    "            48: \"sandwich\",\n",
    "            49: \"orange\",\n",
    "            50: \"broccoli\",\n",
    "            51: \"carrot\",\n",
    "            52: \"hot dog\",\n",
    "            53: \"pizza\",\n",
    "            54: \"donut\",\n",
    "            55: \"cake\",\n",
    "            56: \"chair\",\n",
    "            57: \"couch\",\n",
    "            58: \"potted plant\",\n",
    "            59: \"bed\",\n",
    "            60: \"dining table\",\n",
    "            61: \"toilet\",\n",
    "            62: \"tv\",\n",
    "            63: \"laptop\",\n",
    "            64: \"mouse\",\n",
    "            65: \"remote\",\n",
    "            66: \"keyboard\",\n",
    "            67: \"cell phone\",\n",
    "            68: \"microwave\",\n",
    "            69: \"oven\",\n",
    "            70: \"toaster\",\n",
    "            71: \"sink\",\n",
    "            72: \"refrigerator\",\n",
    "            73: \"book\",\n",
    "            74: \"clock\",\n",
    "            75: \"vase\",\n",
    "            76: \"scissors\",\n",
    "            77: \"teddy bear\",\n",
    "            78: \"hair drier\",\n",
    "            79: \"toothbrush\",\n",
    "        },\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Once the training is complete, the output directory looks like this:\n",
    "```\n",
    "out/my_experiment\n",
    "├── checkpoints\n",
    "│   ├── best.ckpt\n",
    "│   └── last.ckpt\n",
    "├── events.out.tfevents.1764251158.ef9b159fe4b8.273.0\n",
    "├── exported_models\n",
    "│   ├── exported_best.pt\n",
    "│   └── exported_last.pt\n",
    "└── train.log\n",
    "```\n",
    "\n",
    "The best model checkpoint is saved to `out/my_experiment/exported_models/exported_best.pt`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Predict with your own LTDETR weights\n",
    "\n",
    "### Download an example image\n",
    "\n",
    "Download an example image for inference with the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O image.jpg http://images.cocodataset.org/val2017/000000577932.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Then load the model with LightlyTrain's `load_model` function. This will automatically download the model weights and load the model.\n",
    "\n",
    "### Load the model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = lightly_train.load_model(\"out/my_experiment/exported_models/exported_best.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Predict the objects\n",
    "\n",
    "Then run `model.predict` on the image. The method accepts file paths, URLs,\n",
    "PIL Images, or tensors as input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = model.predict(\"image.jpg\")\n",
    "\n",
    "results[\"labels\"]  # Class labels, tensor of shape (num_boxes,)\n",
    "results[\"bboxes\"]  # Bounding boxes in (xmin, ymin, xmax, ymax) absolute pixel\n",
    "# coordinates of the original image. Tensor of shape (num_boxes, 4).\n",
    "results[\"scores\"]  # Confidence scores, tensor of shape (num_boxes,)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Visualize the results\n",
    "\n",
    "Finally, we visualize the image and results to check what objects were detected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torchvision.io import read_image\n",
    "from torchvision.utils import draw_bounding_boxes\n",
    "\n",
    "image = read_image(\"image.jpg\")\n",
    "image_with_boxes = draw_bounding_boxes(\n",
    "    image,\n",
    "    boxes=results[\"bboxes\"],\n",
    "    labels=[model.classes[label.item()] for label in results[\"labels\"]],\n",
    ")\n",
    "plt.imshow(image_with_boxes.permute(1, 2, 0))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## Benchmark your checkpoint\n",
    "\n",
    "You can use the `benchmark_object_detection` (in beta) command to measure the inference performance of an object detection model on a validation dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = lightly_train.benchmark_object_detection(\n",
    "    out=\"out/my_benchmark\",\n",
    "    dataset_name=\"coco128\",  # Human-readable name shown in the report.\n",
    "    model=\"out/my_experiment/exported_models/exported_best.pt\",\n",
    "    data={\n",
    "        \"path\": \"coco128_yolo\",\n",
    "        \"train\": \"images/train2017\",\n",
    "        \"val\": \"images/val2017\",\n",
    "        \"names\": {\n",
    "            0: \"person\",\n",
    "            1: \"bicycle\",\n",
    "            2: \"car\",\n",
    "            3: \"motorcycle\",\n",
    "            4: \"airplane\",\n",
    "            5: \"bus\",\n",
    "            6: \"train\",\n",
    "            7: \"truck\",\n",
    "            8: \"boat\",\n",
    "            9: \"traffic light\",\n",
    "            10: \"fire hydrant\",\n",
    "            11: \"stop sign\",\n",
    "            12: \"parking meter\",\n",
    "            13: \"bench\",\n",
    "            14: \"bird\",\n",
    "            15: \"cat\",\n",
    "            16: \"dog\",\n",
    "            17: \"horse\",\n",
    "            18: \"sheep\",\n",
    "            19: \"cow\",\n",
    "            20: \"elephant\",\n",
    "            21: \"bear\",\n",
    "            22: \"zebra\",\n",
    "            23: \"giraffe\",\n",
    "            24: \"backpack\",\n",
    "            25: \"umbrella\",\n",
    "            26: \"handbag\",\n",
    "            27: \"tie\",\n",
    "            28: \"suitcase\",\n",
    "            29: \"frisbee\",\n",
    "            30: \"skis\",\n",
    "            31: \"snowboard\",\n",
    "            32: \"sports ball\",\n",
    "            33: \"kite\",\n",
    "            34: \"baseball bat\",\n",
    "            35: \"baseball glove\",\n",
    "            36: \"skateboard\",\n",
    "            37: \"surfboard\",\n",
    "            38: \"tennis racket\",\n",
    "            39: \"bottle\",\n",
    "            40: \"wine glass\",\n",
    "            41: \"cup\",\n",
    "            42: \"fork\",\n",
    "            43: \"knife\",\n",
    "            44: \"spoon\",\n",
    "            45: \"bowl\",\n",
    "            46: \"banana\",\n",
    "            47: \"apple\",\n",
    "            48: \"sandwich\",\n",
    "            49: \"orange\",\n",
    "            50: \"broccoli\",\n",
    "            51: \"carrot\",\n",
    "            52: \"hot dog\",\n",
    "            53: \"pizza\",\n",
    "            54: \"donut\",\n",
    "            55: \"cake\",\n",
    "            56: \"chair\",\n",
    "            57: \"couch\",\n",
    "            58: \"potted plant\",\n",
    "            59: \"bed\",\n",
    "            60: \"dining table\",\n",
    "            61: \"toilet\",\n",
    "            62: \"tv\",\n",
    "            63: \"laptop\",\n",
    "            64: \"mouse\",\n",
    "            65: \"remote\",\n",
    "            66: \"keyboard\",\n",
    "            67: \"cell phone\",\n",
    "            68: \"microwave\",\n",
    "            69: \"oven\",\n",
    "            70: \"toaster\",\n",
    "            71: \"sink\",\n",
    "            72: \"refrigerator\",\n",
    "            73: \"book\",\n",
    "            74: \"clock\",\n",
    "            75: \"vase\",\n",
    "            76: \"scissors\",\n",
    "            77: \"teddy bear\",\n",
    "            78: \"hair drier\",\n",
    "            79: \"toothbrush\",\n",
    "        },\n",
    "    },\n",
    ")\n",
    "\n",
    "result.print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "The command returns a `BenchmarkResult` instance and writes two files to the out directory:\n",
    "\n",
    "```\n",
    "out/my_benchmark\n",
    "├── benchmark_results.json\n",
    "└── benchmark_summary.md\n",
    "```\n",
    "\n",
    "* `benchmark_results.json`: the full result as JSON.\n",
    "* `benchmark_summary.md`: a human-readable Markdown report.\n",
    "\n",
    "The report (also available via `result.to_markdown()`) contains the run configuration, device info, performance metrics, and a throughput & latency table."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "* [Object Detection Documentation](https://docs.lightly.ai/train/stable/object_detection.html): If you want to learn more about object detection with LightlyTrain.\n",
    "* [Object Detection Model Export](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/object_detection.ipynb): If you want to learn more about export the model to ONNX and TensorRT format."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}