autodistill / autodistill-seggpt

SegGPT module for use with Autodistill

Autodistill SegGPT Module

This repository contains the code supporting the SegGPT base model for use with Autodistill.

SegGPT is a transformer-based, few-shot semantic segmentation model developed by BAAI Vision.

This model performs well on task-specific segmentation tasks when given a few labeled images from which to learn features about the objects you want to identify.

Read the full Autodistill documentation.

Read the SegGPT Autodistill documentation.

Installation

To use SegGPT with Autodistill, you need to install the following dependency:

pip3 install autodistill-seggpt

About SegGPT

SegGPT performs "in-context" segmentation. This means it requires a handful of pre-labelled "context" images.

You will need some labeled images to use SegGPT. Don't have any labeled images? Check out Roboflow Annotate, a feature-rich annotation tool from which you can export data for use with Autodistill.

Quickstart

from autodistill_seggpt import SegGPT, FewShotOntology

base_model = SegGPT(
    ontology=FewShotOntology(supervision_dataset)
)

base_model.label("./unlabelled-photos", extension=".jpg")

How to load data from Roboflow

Labelling and importing images is easy!

You can use Roboflow Annotate to label a few images (1-3 should work fine). For your Project Type, make sure to pick Instance Segmentation, as you will be labelling with polygons.

Once you have labelled your images, you can press Generate > Generate New Version. You can use all the default options--no Augmentations are necessary.

Once your dataset version is generated, you can press Export > Continue.

Then you will get some download code to copy. It should look something like this:

!pip install roboflow

from roboflow import Roboflow
rf = Roboflow(api_key="ABCDEFG")
project = rf.workspace("lorem-ipsum").project("dolor-sit-amet")
dataset = project.version(1).download("yolov8")

Note: if you are not using a notebook environment, you should remove !pip install roboflow from your code, and run pip install roboflow in your terminal instead.

To import your dataset into Autodistill, run the following:

import supervision as sv

supervision_dataset = sv.DetectionDataset.from_yolo(
    images_directory_path=f"{dataset.location}/train/images",
    annotations_directory_path=f"{dataset.location}/train/labels",
    data_yaml_path=f"{dataset.location}/data.yaml",
    force_masks=True
)

License

The code in this repository is licensed under an MIT license.

See the SegGPT repository for more information on the SegGPT license.

🏆 Contributing

We love your input! Please see the core Autodistill contributing guide to get started. Thank you 🙏 to all our contributors!

autodistill-seggpt's Issues

CUDA out of memory

I'm using this library in a server that has 4 GPU's with 16 GB each GPU. When I set the number to pre-annotated images to more than 8 the one GPU get's without memory making the process to crash.

Is this possible to run the process across multiple GPU's? If not what it's the optimal size from the annotated images that need to be used? Cause I've been resizing my images from 960 down to 448 and even with the low resolution I get the memory error. In which type of GPU is recommended to use this tool?

Memory Leak? CUDA memory finishes quickly when using more than 50 pre-annotated data

Even when using 150 pre annotated images to label 10 additional images, CUDA memory finishes on a 3090 GPU.

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 31.36 GiB (GPU 0; 24.00 GiB total capacity; 49.21 GiB already allocated; 0 bytes free; 49.35 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

It stops working after about 50 pre annotated images due to this error.

No module named 'seggpt_engine'

Traceback (most recent call last):
File "/home/smiv/Temp2018/ZWP/project2/01-yolo/autodistill/s.py", line 1, in
from autodistill_seggpt import SegGPT, FewShotOntology
File "/home/smiv/anaconda3/envs/zwp_autodistill/lib/python3.9/site-packages/autodistill_seggpt/init.py", line 6, in
from .seggpt import SegGPT
File "/home/smiv/anaconda3/envs/zwp_autodistill/lib/python3.9/site-packages/autodistill_seggpt/seggpt.py", line 63, in
from seggpt_engine import run_one_image
ModuleNotFoundError: No module named 'seggpt_engine'

Ability to use different GPU in cluster in seggpt.py

Currently the code is hard coded to use device:0 with device = ("cuda" if torch.cuda.is_available() else "cpu"). It would be a great enhancement to also be able to set an environment variable with the GPU ID that we would want to use when there are multiple GPUs available.

I attempted to set CUDA_VISIBLE_DEVICES=1 to see if it would use the next GPU, but it did not.

So maybe something like in a terminal the user could set:
$export CUDA_DEVICE=1

And in seggpt.py:
if os.getenv("CUDA_DEVICE"):
device = torch.device(f"cuda:{os.getenv('CUDA_DEVICE')}" if torch.cuda.is_available() else "cpu")
else:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

`seggpt.check_dependencies` permanently changes working directory

This line in seggpt.check_dependencies changes the working directory of the Python process, but never reverts to the original working directory once the dependency check is complete. This lack of revert causes paths to be incorrectly resolved in programs that import seggpt.

autodistill-seggpt/autodistill_seggpt/seggpt.py

Line 16 in ad60a7a

os.chdir(autodistill_dir)