CSP-SAM: CNN-Enhanced and Self-Prompting SAM for Ultrasound Anatomical Structure Segmentation

这是一个基于 SAM 的医学图像分割框架，通过 CNN 增强和自提示机制，专为超声解剖结构分割设计。

🚀 Quick Start

# Create conda environment
conda env create -f environment.yml
conda activate medsam

# Training
python train_refine_de.py --data_path /path/to/data --sam_checkpoint /path/to/checkpoint.pth

# Testing
python test_refine_de.py --data_path /path/to/data --sam_checkpoint /path/to/checkpoint.pth

📁 Project Structure

CSP-SAM/
├── train_refine_de.py          # Training script
├── test_refine_de.py           # Testing script
├── DataLoader.py               # Dataset loading utilities
├── utils.py                    # Utility functions
├── metrics.py                  # Segmentation metrics
├── iou_dice.py                 # IoU and Dice calculation
├── refine.py                   # Refinement modules
├── segment_anything/           # Original SAM implementation
├── segment_anything_cnn_refine_decoder_2/  # CNN-enhanced SAM
├── FastGeodis/                 # Fast Geodesic distance
├── scripts/                    # Utility scripts
├── app.ipynb                   # Demo notebook
├── predictor_example.ipynb     # Predictor example
└── environment.yml             # Conda environment

📥 Downloads

Datasets

Dataset	Link
CAMUS	Google Drive
Cradiac	Google Drive
Dynamic	Google Drive

Pretrained Weights

Model	Link
SAM-Med2D	Google Drive

Note: Create a pretrain_model/ folder and place the weights file there.

📁 Dataset Format

data/
└── CAMUS/
    ├── CAMUS_train.json         # Training split
    ├── CAMUS_test.json          # Test split
    ├── images/                  # Image files
    └── masks/                   # Ground truth masks

JSON format:

{
    "/path/to/images/img001.png": ["/path/to/masks/mask001.png"],
    "/path/to/images/img002.png": ["/path/to/masks/mask002_1.png", "/path/to/masks/mask002_2.png"]
}

📤 Outputs

workdir/
├── models/
│   └── run_name/
│       └── epoch*_sam.pth       # Checkpoints
└── logs/
    └── run_name_*.log           # Training logs

📄 License

MIT License. See LICENSE for details.