LightlyTrain for fine-tuning DINO foundation models

[xiaoyachong]

I mimicked the mlex_dlsia_segmentation_prototype repository to create this one. LightlyTrain is a useful package for fine-tuning foundation models for instance, semantic, and panoptic segmentation. I utilized it for semantic segmentation in our segmentation application.

To fine-tune a LightlyTrain model, we first need to convert the Tiled data into a temporary data folder with the following structure:

my_data_dir
├── train
│   ├── images
│   │   ├── image0.jpg
│   │   └── image1.jpg
│   └── masks
│       ├── image0.png
│       └── image1.png
└── val
    ├── images
    │   ├── image2.jpg
    │   └── image3.jpg
    └── masks
        ├── image2.png
        └── image3.png

After the data is prepared, we can call lightly_train.train_semantic_segmentation to fine-tune the model:

import lightly_train

if __name__ == "__main__":
    lightly_train.train_semantic_segmentation(
        out="out/my_experiment",
        model="dinov3/vits16-eomt",
        data={
            "train": {
                "images": "my_data_dir/train/images",
                "masks": "my_data_dir/train/masks",
            },
            "val": {
                "images": "my_data_dir/val/images",
                "masks": "my_data_dir/val/masks",
            },
            "classes": {
                0: "sand",
                1: "air",
                2: "background",
                255: "unknown",
            },
            "ignore_classes": [255],
        },
    )

Once fine-tuning is complete, the fine-tuned model checkpoint is registered to MLflow.

Companion PRs:
mlex_highres: mlexchange/mlex_highres_segmentation#219
mlex_tomo: mlexchange/mlex_tomo_framework#18

[taxe10]

PR looks great! I left some minor comments. We should also add some unit tests to the repo.

[xiaoyachong]

PR looks great! I left some minor comments. We should also add some unit tests to the repo.

Thanks for your comments! I just addressed them and added pytest.