/finetune-anything

Fine-tune SAM (Segment Anything Model) for class-aware computer vision tasks in specific scenarios

Primary LanguagePythonMIT LicenseMIT

Introduction

The Segment Anything Model (SAM) has revolutionized computer vision. Relying on fine-tuning of SAM will solve a large number of basic computer vision tasks. We are designing a class-aware one-stage tool for training fine-tuning models based on SAM.

You need to supply the datasets for your tasks and the supported task name, this tool will help you to get a finetuned model for your task.

Design

Finetune-Anything further encapsulates the three parts of the original SAM, i.e., Image Encoder Adapter, Prompt Encoder Adapter, and Mask Decoder Adatper. We will support the base extend-SAM model for each task. Users also could design your own customized modules in each adapter, check details in How_to_use. For example, MaskDecoder is encapsulated as MaskDecoderAdapter. The current MaskDecoderAdatper contains two parts, DecoderNeck and DecoderHead.

Supported Tasks

  • Semantic Segmentation
    • train
    • eval
    • test
  • Matting
  • Instance Segmentation
  • Detection

Supported Datasets

  • TorchVOCSegmentation
  • BaseSemantic
  • BaseInstance
  • BaseMatting

Deploy

  • Onnx export

Usage

finetune-anything(FA) supports the entire training process of SAM model fine-tuning, including the modification of the model structure, as well as the model training, verification, and testing processes. For details, check the How_to_use, the Quick Start gives an example of quickly using FA to train a custom semantic segmentation model.

Quick Start

Install

  • Step1
git clone https://github.com/ziqi-jin/finetune-anything.git
cd finetune-anything
pip install -r requirements.txt

  • Step2 Download the SAM weights from SAM repository

  • Step3 Modify the contents of yaml file for the specific task in /config, e.g., ckpt_path, model_type ...

Train

CUDA_VISIBLE_DEVICES=${your GPU number} python train.py --task_name semantic_seg