{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Instance Segmentation with LTDETRv2\n",
    "\n",
    "This notebook demonstrates how to use LightlyTrain for instance segmentation with our state-of-the-art LTDETRv2 model and publicly released weights trained on the [COCO](https://arxiv.org/abs/1612.03716) dataset.\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/ltdetr_instance_segmentation.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"
   ]
  },
  {
   "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 instance segmentation model\n",
    "\n",
    "Training your own LTDETRv2 instance segmentation model is straightforward with LightlyTrain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "### Download dataset\n",
    "\n",
    "First download a dataset in YOLO segmentation format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O coco128-seg.zip https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128-seg.zip && unzip -q coco128-seg.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "The example archive contains a single image directory. Move a small subset into a validation split so that the tutorial demonstrates both training and validation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "dataset_dir = Path(\"coco128-seg\")\n",
    "val_images_dir = dataset_dir / \"images/val2017\"\n",
    "val_labels_dir = dataset_dir / \"labels/val2017\"\n",
    "val_images_dir.mkdir(parents=True, exist_ok=True)\n",
    "val_labels_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "if not any(val_images_dir.iterdir()):\n",
    "    train_images_dir = dataset_dir / \"images/train2017\"\n",
    "    train_labels_dir = dataset_dir / \"labels/train2017\"\n",
    "    for image_path in sorted(train_images_dir.glob(\"*.jpg\"))[-16:]:\n",
    "        image_path.rename(val_images_dir / image_path.name)\n",
    "        label_path = train_labels_dir / image_path.with_suffix(\".txt\").name\n",
    "        if label_path.exists():\n",
    "            label_path.rename(val_labels_dir / label_path.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Then start training with `train_instance_segmentation`. LightlyTrain automatically configures the model and augmentations, while still allowing every setting to be customized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "lightly_train.train_instance_segmentation(\n",
    "    out=\"out/my_experiment\",\n",
    "    model=\"ltdetrv2-seg-s-coco\",\n",
    "    steps=100,  # Small number of steps for demonstration.\n",
    "    batch_size=4,  # Small batch size for demonstration.\n",
    "    data={\n",
    "        \"path\": \"coco128-seg\",\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": "11",
   "metadata": {},
   "source": [
    "Once training completes, the final model checkpoint is saved in `out/my_experiment/exported_models/exported_last.pt`. The best model according to the validation mask mAP is saved in `out/my_experiment/exported_models/exported_best.pt`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Predict with your own LTDETRv2 weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O image.jpg http://images.cocodataset.org/val2017/000000039769.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Load the model weights\n",
    "\n",
    "Load the best model exported during training with LightlyTrain's `load_model` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = lightly_train.load_model(\"out/my_experiment/exported_models/exported_best.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "### Predict and visualize the instances\n",
    "\n",
    "Run prediction on the example image with your trained model and visualize its masks, bounding boxes, and class labels:"
   ]
  },
  {
   "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, draw_segmentation_masks\n",
    "\n",
    "prediction = model.predict(\"image.jpg\")\n",
    "\n",
    "image = read_image(\"image.jpg\")\n",
    "image_with_masks = draw_segmentation_masks(image, masks=prediction[\"masks\"], alpha=0.6)\n",
    "image_with_instances = draw_bounding_boxes(\n",
    "    image_with_masks,\n",
    "    boxes=prediction[\"bboxes\"],\n",
    "    labels=[model.classes[label.item()] for label in prediction[\"labels\"]],\n",
    ")\n",
    "plt.imshow(image_with_instances.permute(1, 2, 0))\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "- [Instance Segmentation Documentation](https://docs.lightly.ai/train/stable/instance_segmentation.html): Learn more about training, inference, data formats, and available models.\n",
    "- [LTDETRv2 Instance Segmentation Model Export](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/ltdetr_instance_segmentation_export.ipynb): Export an LTDETRv2 instance segmentation model to ONNX and TensorRT."
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}