Instance Segmentation with LTDETRv2

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 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 LTDETRv2 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 training with train_instance_segmentation. LightlyTrain automatically configures the model and augmentations, while still allowing every setting to be customized.

import lightly_train

lightly_train.train_instance_segmentation(
    out="out/my_experiment",
    model="ltdetrv2-seg-s-coco",
    steps=100,  # Small number of steps for demonstration.
    batch_size=4,  # Small batch size for demonstration.
    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 LTDETRv2 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 with your trained model and visualize its masks, bounding boxes, and class labels:

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

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

image = read_image("image.jpg")
image_with_masks = draw_segmentation_masks(image, masks=prediction["masks"], alpha=0.6)
image_with_instances = draw_bounding_boxes(
    image_with_masks,
    boxes=prediction["bboxes"],
    labels=[model.classes[label.item()] for label in prediction["labels"]],
)
plt.imshow(image_with_instances.permute(1, 2, 0))
plt.axis("off")
plt.show()

Next Steps