{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Panoptic Segmentation with DINOv3 EoMT\n",
    "\n",
    "This notebook demonstrates how to use LightlyTrain for panoptic segmentation with our\n",
    "state-of-the-art [EoMT](https://arxiv.org/abs/2503.19108) model built on [DINOv3](https://github.com/facebookresearch/dinov3)\n",
    "backbones, with our publicly released weights trained on the [COCO](https://arxiv.org/abs/1612.03716)\n",
    "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/eomt_panoptic_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": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install lightly-train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "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": "3",
   "metadata": {},
   "source": [
    "## Train a panoptic segmentation model\n",
    "\n",
    "Training your own panoptic segmentation model is straightforward with LightlyTrain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "### Prepare Data\n",
    "\n",
    "LightlyTrain supports panoptic segmentation datasets in COCO format.\n",
    "Every image must have a corresponding mask image that encodes the segmentation class\n",
    "and segment ID for each pixel. The dataset must also include COCO-style JSON annotation\n",
    "files.\n",
    "\n",
    "Your dataset directory must be organized like this:\n",
    "\n",
    "```text\n",
    "my_data_dir/\n",
    "├── images\n",
    "│   ├── train\n",
    "│   │   ├── image1.jpg\n",
    "│   │   └── ...\n",
    "│   └── val\n",
    "│       ├── image2.jpg\n",
    "│       └── ...\n",
    "└── annotations\n",
    "    ├── train\n",
    "    │   ├── image1.png\n",
    "    │   └── ...\n",
    "    ├── train.json\n",
    "    ├── val\n",
    "    │   ├── image2.png\n",
    "    │   └── ...\n",
    "    └── val.json\n",
    "```\n",
    "\n",
    "The directory names don't matter as long as you provide the correct paths in the\n",
    "training function.\n",
    "\n",
    "You can download an example COCO dataset from here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/lightly-ai/coco128_panoptic/releases/download/v0.0.1/coco128_panoptic.zip && unzip -q coco128_panoptic.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "The dataset looks like this after the download completes:\n",
    "\n",
    "```\n",
    "coco128_panoptic\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",
    "└── annotations\n",
    "    ├── panoptic_train2017\n",
    "    │   ├── 000000000009.png\n",
    "    │   ├── 000000000025.png\n",
    "    │   ├── ...\n",
    "    │   └── 000000000659.png\n",
    "    ├── panoptic_train2017.json\n",
    "    ├── panoptic_val2017\n",
    "    │   ├──  000000000139.png\n",
    "    │   ├──  000000000285.png\n",
    "    │   ├──  ...\n",
    "    │   └──  000000013201.png\n",
    "    └── panoptic_val2017.json\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "### Start Training\n",
    "\n",
    "Then start the training with the `train_panoptic_segmentation` function. You can specify various training parameters such as the model architecture, number of training steps, batch size, learning rate, and more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "lightly_train.train_panoptic_segmentation(\n",
    "    out=\"out/my_experiment\",\n",
    "    model=\"dinov3/vits16-eomt-panoptic-coco\",\n",
    "    steps=100,  # Small number of steps for demonstration, default is 90_000.\n",
    "    batch_size=4,  # Small batch size for demonstration, default is 16.\n",
    "    data={\n",
    "        \"train\": {\n",
    "            \"images\": \"coco128_panoptic/images/train2017\",  # Path to train images\n",
    "            \"masks\": \"coco128_panoptic/annotations/panoptic_train2017\",  # Path to train mask images\n",
    "            \"annotations\": \"coco128_panoptic/annotations/panoptic_train2017.json\",  # Path to train COCO-style annotations\n",
    "        },\n",
    "        \"val\": {\n",
    "            \"images\": \"coco128_panoptic/images/val2017\",  # Path to val images\n",
    "            \"masks\": \"coco128_panoptic/annotations/panoptic_val2017\",  # Path to val mask images\n",
    "            \"annotations\": \"coco128_panoptic/annotations/panoptic_val2017.json\",  # Path to val COCO-style annotations\n",
    "        },\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Once training completes, the final model checkpoint is saved in `out/my_experiment/exported_models/exported_last.pt`.\n",
    "The best model according to the validation panoptic quality is\n",
    "saved in `out/my_experiment/exported_models/exported_best.pt`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Predict with your own EoMT weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O image.jpg http://images.cocodataset.org/val2017/000000070254.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "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": "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 and visualize the panoptic segmentation\n",
    "\n",
    "Run prediction on the example image again, this time using the weights from your own training run, and visualize the panoptic masks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torchvision.io import read_image\n",
    "from torchvision.utils import draw_segmentation_masks\n",
    "\n",
    "image = read_image(\"image.jpg\")\n",
    "results = model.predict(\"image.jpg\")\n",
    "\n",
    "masks = results[\"masks\"]\n",
    "segment_ids = results[\"segment_ids\"]\n",
    "masks_bool = torch.stack(\n",
    "    [masks[..., 1] == -1] + [masks[..., 1] == segment_id for segment_id in segment_ids]\n",
    ")\n",
    "colors = [(0, 0, 0)] + [\n",
    "    [int(color * 255) for color in plt.cm.tab20c(i / len(segment_ids))[:3]]\n",
    "    for i in range(len(segment_ids))\n",
    "]\n",
    "image_with_masks = draw_segmentation_masks(image, masks_bool, colors=colors, alpha=1.0)\n",
    "plt.imshow(image_with_masks.permute(1, 2, 0))\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}