Semantic Segmentation with DINOv3 EoMT

This notebook demonstrates how to use LightlyTrain for semantic segmentation with our state-of-the-art EoMT model built on DINOv3 backbones, with our publicly released weights trained on the COCO-Stuff and Cityscapes dataset.

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.

Installation

LightlyTrain can be installed directly via pip:

!pip install lightly-train

Important: LightlyTrain is officially supported on

• Linux: CPU or CUDA

• MacOS: CPU only

• Windows (experimental): CPU or CUDA

We are planning to support MPS for MacOS.

Check the installation instructions for more details on installation.

Train a semantic segmentation model

Training your own semantic segmentation model is straightforward with LightlyTrain.

Download and prepare a dataset

First download a small dataset in YOLO segmentation format. The following cells create a validation split and convert its polygon annotations into semantic masks.

!wget -O coco128-seg.zip https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128-seg.zip && unzip -q coco128-seg.zip

from pathlib import Path

from PIL import Image, ImageDraw

dataset_dir = Path("coco128-seg")
val_images_dir = dataset_dir / "images/val2017"
val_labels_dir = dataset_dir / "labels/val2017"
val_images_dir.mkdir(parents=True, exist_ok=True)
val_labels_dir.mkdir(parents=True, exist_ok=True)

if not any(val_images_dir.iterdir()):
    train_images_dir = dataset_dir / "images/train2017"
    train_labels_dir = dataset_dir / "labels/train2017"
    for image_path in sorted(train_images_dir.glob("*.jpg"))[-16:]:
        image_path.rename(val_images_dir / image_path.name)
        label_path = train_labels_dir / image_path.with_suffix(".txt").name
        if label_path.exists():
            label_path.rename(val_labels_dir / label_path.name)

coco_class_names = (
    "person",
    "bicycle",
    "car",
    "motorcycle",
    "airplane",
    "bus",
    "train",
    "truck",
    "boat",
    "traffic light",
    "fire hydrant",
    "stop sign",
    "parking meter",
    "bench",
    "bird",
    "cat",
    "dog",
    "horse",
    "sheep",
    "cow",
    "elephant",
    "bear",
    "zebra",
    "giraffe",
    "backpack",
    "umbrella",
    "handbag",
    "tie",
    "suitcase",
    "frisbee",
    "skis",
    "snowboard",
    "sports ball",
    "kite",
    "baseball bat",
    "baseball glove",
    "skateboard",
    "surfboard",
    "tennis racket",
    "bottle",
    "wine glass",
    "cup",
    "fork",
    "knife",
    "spoon",
    "bowl",
    "banana",
    "apple",
    "sandwich",
    "orange",
    "broccoli",
    "carrot",
    "hot dog",
    "pizza",
    "donut",
    "cake",
    "chair",
    "couch",
    "potted plant",
    "bed",
    "dining table",
    "toilet",
    "tv",
    "laptop",
    "mouse",
    "remote",
    "keyboard",
    "cell phone",
    "microwave",
    "oven",
    "toaster",
    "sink",
    "refrigerator",
    "book",
    "clock",
    "vase",
    "scissors",
    "teddy bear",
    "hair drier",
    "toothbrush",
)
class_ids = set()

for split in ("train2017", "val2017"):
    images_dir = dataset_dir / "images" / split
    labels_dir = dataset_dir / "labels" / split
    masks_dir = dataset_dir / "masks" / split
    masks_dir.mkdir(parents=True, exist_ok=True)
    for image_path in images_dir.glob("*.jpg"):
        with Image.open(image_path) as image:
            width, height = image.size
        mask = Image.new("L", (width, height), 0)
        draw = ImageDraw.Draw(mask)
        label_path = labels_dir / image_path.with_suffix(".txt").name
        lines = label_path.read_text().splitlines() if label_path.exists() else ()
        for line in lines:
            values = [float(value) for value in line.split()]
            class_id = int(values[0])
            class_ids.add(class_id)
            polygon = [
                (int(x * width), int(y * height))
                for x, y in zip(values[1::2], values[2::2])
            ]
            draw.polygon(polygon, fill=class_id + 1)
        mask.save(masks_dir / image_path.with_suffix(".png").name)

classes = {0: "background"}
classes.update(
    {class_id + 1: coco_class_names[class_id] for class_id in sorted(class_ids)}
)

Start training

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.

import lightly_train

lightly_train.train_semantic_segmentation(
    out="out/my_experiment",
    model="dinov3/vits16-eomt-coco",
    steps=100,  # Small number of steps for demonstration, default is 90_000.
    batch_size=4,  # Small batch size for demonstration, default is 16.
    data={
        "train": {
            "images": "coco128-seg/images/train2017",
            "masks": "coco128-seg/masks/train2017",
        },
        "val": {
            "images": "coco128-seg/images/val2017",
            "masks": "coco128-seg/masks/val2017",
        },
        "classes": classes,
    },
)

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.

Predict with your own EoMT weights

!wget -O cat.jpg http://images.cocodataset.org/val2017/000000039769.jpg

Load the model weights

Load the best model exported during training with LightlyTrain’s load_model function:

model = lightly_train.load_model("out/my_experiment/exported_models/exported_best.pt")

Predict and visualize the semantic mask

Run prediction on the example image again, this time using the weights from your own training run, and visualize the semantic mask:

import matplotlib.pyplot as plt
import torch
from torchvision.io import read_image
from torchvision.utils import draw_segmentation_masks

masks = model.predict("cat.jpg")

image = read_image("cat.jpg")
masks = torch.stack([masks == class_id for class_id in masks.unique()])
image_with_masks = draw_segmentation_masks(image, masks, alpha=0.6)
plt.imshow(image_with_masks.permute(1, 2, 0))
plt.axis("off")
plt.show()