Update README.md

c8843a6 verified about 2 months ago

3.24 kB

license: apache-2.0
tags:
  - medical-imaging
  - image-segmentation
  - semi-supervised-learning
  - sam3
  - angiography
  - mean-teacher
  - pytorch

SMART: Semi-supervised Medical Adaptive vessel Representation Toolkit

This repository hosts the official pre-trained and fine-tuned model checkpoints for the paper SMART (arXiv:2603.00881).

SMART is a semi-supervised video vessel segmentation framework built on SAM3 (Segment Anything Model 3). It features a Mean Teacher architecture and text prompt support, specifically designed for efficient coronary angiography vessel segmentation.

🗂️ Model Weights Overview

We provide all necessary weights to reproduce our experiments, from the original baselines to our final semi-supervised checkpoints:

Filename	Size	Description
`sam2.1_hiera_large.pt`	898 MB	Original SAM 2.1 Hiera Large baseline weights.
`sam3_original.pt`	3.45 GB	Original SAM 3 baseline weights.
`sam3_1p_finetune_checkpoint_100.pt`	10.1 GB	Supervised fine-tuning checkpoint trained using only 1% labeled data (100 epochs).
`semi_sam3_5labeled_checkpoint_final.pt`	10.6 GB	Final SMART checkpoint trained via semi-supervised learning (Mean Teacher) with 5% labeled data.
`bpe_simple_vocab_16e6.txt.gz`	1.36 MB	BPE vocabulary file required for the text prompt tokenizer.

🚀 How to Use

You can easily download these weights using the huggingface_hub library and integrate them directly into the SMART training/inference pipeline.

1. Install Dependencies

pip install huggingface_hub torch torchvision

2. Download and Load Checkpoints

Here is an example of how to download the final SMART checkpoint and load it into your PyTorch environment:

from huggingface_hub import hf_hub_download
import torch

# 1. Download the final semi-supervised checkpoint
ckpt_path = hf_hub_download(
    repo_id="ly17/TC-SemiSAM-checkpoints", 
    filename="semi_sam3_5labeled_checkpoint_final.pt"
)
print(f"Weights downloaded to: {ckpt_path}")

# 2. Load the state dict
# state_dict = torch.load(ckpt_path, map_location="cpu")
# model.load_state_dict(state_dict)

3. Text Prompt Setup

When running inference with our SMART model, please ensure you use the following default text prompt as specified in our methodology:

TEXT_PROMPT = "Please segment the blood vessels"

(Note: The model expects a dataset resolution of 512×512, which is resized to a SAM3 input resolution of 1008×1008 with normalization range [-1, 1].)

⚠️ Clinical Disclaimer

These models are released for research purposes only. They are not intended for direct clinical decision-making, patient diagnosis, or treatment planning.

While we would love to host the XCA angiography sequences on HF Datasets, they contain sensitive medical imaging data.