Home

Awesome

Lightning Segment-Anything Model

This library allows you to fine-tune the powerful Segment-Anything model from MetaAI for your custom COCO-format datasets. The library is built on top of Lightning AI's Fabric framework, providing an efficient and easy-to-use implementation for achieving state-of-the-art instance segmentation results.

This repo is an experiment; A PoC to see if fine-tuning SAM using bounding boxes as prompts would increase the IoU or improve the quality of the masks in general. One can use a COCO format dataset to fine-tune SAM for a specific task where SAM does not perform well (e.g., segmenting text on documents) and then use that model with interactive prompts just like SAM.

For generating bboxes from text and prompt them to SAM you may check: lang-segment-anything

Features

Results

Using bbox prompt, the table below shows the mask quality of the original and fine-tuned models.

DatasetTypeMean IoUMean F1Epoch
COCO2017Original SAM0.79780.8749
COCO2017Fine-tuned0.80700.88162
TrashCan 1.0Original SAM0.66430.7808
TrashCan 1.0Fine-tuned0.78880.87385
TrashCan 1.0Fine-tuned0.80000.879514
WGISDOriginal SAM0.85360.9178
WGISDFine-tuned0.87320.929819
<details> <summary>Comparison images</summary> <br> TrashCan 1.0 GT | 6 Epochs Fine-tuned SAM | Original SAM

1 2 3

</details>

Installation

git clone https://github.com/luca-medeiros/lightning-sam.git
cd lightning-sam
pip install .

Quick Start

  1. Prepare your custom COCO-format dataset. The dataset should include a JSON file with annotations and an images folder with corresponding images.

  2. Edit src/config.py with your dataset paths.

  3. Run src/train.py.

Notes

Resources

License

This project is licensed same as SAM model.