{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Semantic Segmentation with DINOv3 EoMT\n",
    "\n",
    "This notebook demonstrates how to use LightlyTrain for semantic segmentation with our state-of-the-art [EoMT](https://arxiv.org/abs/2503.19108) model built on [DINOv3](https://github.com/facebookresearch/dinov3) backbones, with our publicly released weights trained on the [COCO-Stuff](https://arxiv.org/abs/1612.03716) and [Cityscapes](https://www.cityscapes-dataset.com/) 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_semantic_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\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 a semantic segmentation model\n",
    "\n",
    "Training your own semantic segmentation model is straightforward with LightlyTrain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "### Download and prepare a dataset\n",
    "\n",
    "First download a small dataset in YOLO segmentation format. The following cells create a validation split and convert its polygon annotations into semantic masks."
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from PIL import Image, ImageDraw\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)\n",
    "\n",
    "coco_class_names = (\n",
    "    \"person\",\n",
    "    \"bicycle\",\n",
    "    \"car\",\n",
    "    \"motorcycle\",\n",
    "    \"airplane\",\n",
    "    \"bus\",\n",
    "    \"train\",\n",
    "    \"truck\",\n",
    "    \"boat\",\n",
    "    \"traffic light\",\n",
    "    \"fire hydrant\",\n",
    "    \"stop sign\",\n",
    "    \"parking meter\",\n",
    "    \"bench\",\n",
    "    \"bird\",\n",
    "    \"cat\",\n",
    "    \"dog\",\n",
    "    \"horse\",\n",
    "    \"sheep\",\n",
    "    \"cow\",\n",
    "    \"elephant\",\n",
    "    \"bear\",\n",
    "    \"zebra\",\n",
    "    \"giraffe\",\n",
    "    \"backpack\",\n",
    "    \"umbrella\",\n",
    "    \"handbag\",\n",
    "    \"tie\",\n",
    "    \"suitcase\",\n",
    "    \"frisbee\",\n",
    "    \"skis\",\n",
    "    \"snowboard\",\n",
    "    \"sports ball\",\n",
    "    \"kite\",\n",
    "    \"baseball bat\",\n",
    "    \"baseball glove\",\n",
    "    \"skateboard\",\n",
    "    \"surfboard\",\n",
    "    \"tennis racket\",\n",
    "    \"bottle\",\n",
    "    \"wine glass\",\n",
    "    \"cup\",\n",
    "    \"fork\",\n",
    "    \"knife\",\n",
    "    \"spoon\",\n",
    "    \"bowl\",\n",
    "    \"banana\",\n",
    "    \"apple\",\n",
    "    \"sandwich\",\n",
    "    \"orange\",\n",
    "    \"broccoli\",\n",
    "    \"carrot\",\n",
    "    \"hot dog\",\n",
    "    \"pizza\",\n",
    "    \"donut\",\n",
    "    \"cake\",\n",
    "    \"chair\",\n",
    "    \"couch\",\n",
    "    \"potted plant\",\n",
    "    \"bed\",\n",
    "    \"dining table\",\n",
    "    \"toilet\",\n",
    "    \"tv\",\n",
    "    \"laptop\",\n",
    "    \"mouse\",\n",
    "    \"remote\",\n",
    "    \"keyboard\",\n",
    "    \"cell phone\",\n",
    "    \"microwave\",\n",
    "    \"oven\",\n",
    "    \"toaster\",\n",
    "    \"sink\",\n",
    "    \"refrigerator\",\n",
    "    \"book\",\n",
    "    \"clock\",\n",
    "    \"vase\",\n",
    "    \"scissors\",\n",
    "    \"teddy bear\",\n",
    "    \"hair drier\",\n",
    "    \"toothbrush\",\n",
    ")\n",
    "class_ids = set()\n",
    "\n",
    "for split in (\"train2017\", \"val2017\"):\n",
    "    images_dir = dataset_dir / \"images\" / split\n",
    "    labels_dir = dataset_dir / \"labels\" / split\n",
    "    masks_dir = dataset_dir / \"masks\" / split\n",
    "    masks_dir.mkdir(parents=True, exist_ok=True)\n",
    "    for image_path in images_dir.glob(\"*.jpg\"):\n",
    "        with Image.open(image_path) as image:\n",
    "            width, height = image.size\n",
    "        mask = Image.new(\"L\", (width, height), 0)\n",
    "        draw = ImageDraw.Draw(mask)\n",
    "        label_path = labels_dir / image_path.with_suffix(\".txt\").name\n",
    "        lines = label_path.read_text().splitlines() if label_path.exists() else ()\n",
    "        for line in lines:\n",
    "            values = [float(value) for value in line.split()]\n",
    "            class_id = int(values[0])\n",
    "            class_ids.add(class_id)\n",
    "            polygon = [\n",
    "                (int(x * width), int(y * height))\n",
    "                for x, y in zip(values[1::2], values[2::2])\n",
    "            ]\n",
    "            draw.polygon(polygon, fill=class_id + 1)\n",
    "        mask.save(masks_dir / image_path.with_suffix(\".png\").name)\n",
    "\n",
    "classes = {0: \"background\"}\n",
    "classes.update(\n",
    "    {class_id + 1: coco_class_names[class_id] for class_id in sorted(class_ids)}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Start training\n",
    "\n",
    "Start training with the `train_semantic_segmentation` function. You only have to specify the output directory, model, and input data; LightlyTrain configures the remaining training parameters and augmentations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightly_train\n",
    "\n",
    "lightly_train.train_semantic_segmentation(\n",
    "    out=\"out/my_experiment\",\n",
    "    model=\"dinov3/vits16-eomt-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-seg/images/train2017\",\n",
    "            \"masks\": \"coco128-seg/masks/train2017\",\n",
    "        },\n",
    "        \"val\": {\n",
    "            \"images\": \"coco128-seg/images/val2017\",\n",
    "            \"masks\": \"coco128-seg/masks/val2017\",\n",
    "        },\n",
    "        \"classes\": classes,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "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 mIoU is saved in `out/my_experiment/exported_models/exported_best.pt`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Predict with your own EoMT weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O cat.jpg http://images.cocodataset.org/val2017/000000039769.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "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": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = lightly_train.load_model(\"out/my_experiment/exported_models/exported_best.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "### Predict and visualize the semantic mask\n",
    "\n",
    "Run prediction on the example image again, this time using the weights from your own training run, and visualize the semantic mask:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "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",
    "masks = model.predict(\"cat.jpg\")\n",
    "\n",
    "image = read_image(\"cat.jpg\")\n",
    "masks = torch.stack([masks == class_id for class_id in masks.unique()])\n",
    "image_with_masks = draw_segmentation_masks(image, masks, alpha=0.6)\n",
    "plt.imshow(image_with_masks.permute(1, 2, 0))\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}