Instance Segmentation with DINOv3 EoMT

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

Open In Colab

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

!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 an instance segmentation model

Training your own instance segmentation model is straightforward with LightlyTrain.

Download dataset

First download a dataset in YOLO segmentation format.

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

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:

from pathlib import Path

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)

Then start the training with the train_instance_segmentation function. You can specify various training parameters such as the model architecture, number of training steps, batch size, learning rate, and more.

import lightly_train

lightly_train.train_instance_segmentation(
    out="out/my_experiment",
    model="dinov3/vits16-eomt-inst-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={
        "path": "coco128-seg",
        "train": "images/train2017",
        "val": "images/val2017",
        "names": {
            0: "person",
            1: "bicycle",
            2: "car",
            3: "motorcycle",
            4: "airplane",
            5: "bus",
            6: "train",
            7: "truck",
            8: "boat",
            9: "traffic light",
            10: "fire hydrant",
            11: "stop sign",
            12: "parking meter",
            13: "bench",
            14: "bird",
            15: "cat",
            16: "dog",
            17: "horse",
            18: "sheep",
            19: "cow",
            20: "elephant",
            21: "bear",
            22: "zebra",
            23: "giraffe",
            24: "backpack",
            25: "umbrella",
            26: "handbag",
            27: "tie",
            28: "suitcase",
            29: "frisbee",
            30: "skis",
            31: "snowboard",
            32: "sports ball",
            33: "kite",
            34: "baseball bat",
            35: "baseball glove",
            36: "skateboard",
            37: "surfboard",
            38: "tennis racket",
            39: "bottle",
            40: "wine glass",
            41: "cup",
            42: "fork",
            43: "knife",
            44: "spoon",
            45: "bowl",
            46: "banana",
            47: "apple",
            48: "sandwich",
            49: "orange",
            50: "broccoli",
            51: "carrot",
            52: "hot dog",
            53: "pizza",
            54: "donut",
            55: "cake",
            56: "chair",
            57: "couch",
            58: "potted plant",
            59: "bed",
            60: "dining table",
            61: "toilet",
            62: "tv",
            63: "laptop",
            64: "mouse",
            65: "remote",
            66: "keyboard",
            67: "cell phone",
            68: "microwave",
            69: "oven",
            70: "toaster",
            71: "sink",
            72: "refrigerator",
            73: "book",
            74: "clock",
            75: "vase",
            76: "scissors",
            77: "teddy bear",
            78: "hair drier",
            79: "toothbrush",
        },
    },
)

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.

Predict with your own EoMT weights

!wget -O image.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 instances

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

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

prediction = model.predict("image.jpg")

image = read_image("image.jpg")
masks = prediction["masks"]
image_with_masks = draw_segmentation_masks(
    image,
    masks=masks,
    alpha=1.0,
)
plt.imshow(image_with_masks.permute(1, 2, 0))
plt.axis("off")
plt.show()