Skip to main content

Neuro-SAM: Foundation Models for Dendrite and Dendritic Spine Segmentation

Project description

Neuro-SAM

Foundation Models from Dendrite and Dendritic Spine Segmentation

python pytorch wandb

This project demonstrates an interactive UI to segment dendrites and dendritic spines. The model of choice is SAMv2 and the framework used is pytorch.

A stack of neural dendrites and dendritic spines

๐Ÿ“ Table of Contents

๐Ÿง  Overview

Neuro-SAM provides an end-to-end pipeline for analyzing neural structures from 3D microscopy data, featuring:

  • Path Tracing: Waypoint-based A* pathfinding
  • Dendrite Segmentation: SAM2-based dendrite segmentation
  • Smart Spine Detection: Multi-view analysis for spine detection
  • Spine Segmentation: Individual spine segmentation using trained SAM2 model

๐Ÿ“ฆ Built With

PyTorch - an open-source machine learning library for Python, widely used for deep learning applications.

Segment Anything Model - a foundation model used for segmentation built by Meta AI.

Weights and Biases - a tool for tracking and visualizing machine learning experiments.

Visual Studio Code - a code editor redefined and optimized for building applications.

FAU High Performance Computing - a high-performance computing cluster at Friedrich-Alexander-Universitรคt Erlangen-Nรผrnberg.

๐Ÿ“ Repository Structure

Neuro-SAM/
โ”œโ”€โ”€ Train-SAMv2/                    # SAM2 training infrastructure
โ”‚   โ”œโ”€โ”€ sam2/                       # SAM2 model implementation
โ”‚   โ”œโ”€โ”€ checkpoints/                # Model checkpoints
โ”‚   โ”œโ”€โ”€ results/                    # Trained model outputs
โ”‚   โ”œโ”€โ”€ utils/                      # Training utilities
โ”‚   โ”œโ”€โ”€ train_dendrites.py         # Dendrite model training
โ”‚   โ””โ”€โ”€ train_spines.py            # Spine model training
โ”œโ”€โ”€ brightest_path_lib/             # Advanced pathfinding algorithms
โ”‚   โ”œโ”€โ”€ algorithm/                  # A* and waypoint search implementations
โ”‚   โ”œโ”€โ”€ cost/                       # Cost function definitions
โ”‚   โ”œโ”€โ”€ heuristic/                  # Heuristic functions
โ”‚   โ”œโ”€โ”€ visualization/              # Path visualization tools
โ”‚   โ””โ”€โ”€ ...
โ”œโ”€โ”€ napari_utils/                   # Napari plugin components
โ”‚   โ”œโ”€โ”€ main_widget.py             # Main interface with anisotropic scaling
โ”‚   โ”œโ”€โ”€ path_tracing_module.py     # Interactive path tracing
โ”‚   โ”œโ”€โ”€ segmentation_module.py     # Dendrite segmentation interface
โ”‚   โ”œโ”€โ”€ spine_detection_module.py  # Spine detection with smart tracking
โ”‚   โ”œโ”€โ”€ spine_segmentation_module.py # Individual spine segmentation
โ”‚   โ””โ”€โ”€ visualization_module.py    # Path management and visualization
โ””โ”€โ”€ neuro_sam_plugin.py            # Main plugin entry point

๐Ÿš€ Installation

Prerequisites

  • Python 3.10+
  • CUDA-compatible GPU (recommended)
  • Conda/Miniconda

Environment Setup

  1. Clone the repository:
git clone https://github.com/nipunarora8/Neuro-SAM.git
cd Neuro-SAM
  1. Create local environment:
conda create -p ./.venv python=3.10 -c conda-forge
conda activate ./.venv
  1. Install dependencies:
pip install uv
uv sync
  1. Download SAM2 checkpoints:
cd Train-SAMv2/checkpoints
bash download_ckpts.sh

๐Ÿ“Š Usage

Quick Start

from neuro_sam_plugin import run_neuro_sam

# Launch with default spacing (94nm x 94nm x 500nm)
viewer = run_neuro_sam(image_path="your_image.tif")

# Launch with custom voxel spacing
viewer = run_neuro_sam(
    image_path="your_image.tif",
    spacing_xyz=(100.0, 100.0, 300.0)  # X, Y, Z spacing in nm
)

Command Line Interface

# Basic usage
python neuro_sam_plugin.py --image_path /path/to/your/image.tif

# Custom spacing
python neuro_sam_plugin.py --image_path image.tif \
    --x-spacing 100.0 --y-spacing 100.0 --z-spacing 300.0

# Load benchmark dataset
python neuro_sam_plugin.py

๐Ÿ”ฌ Workflow

1. Configure Voxel Spacing

Set accurate X, Y, Z voxel spacing in the "Path Tracing" tab for proper anisotropic scaling:

  • Typical two-photon: 94nm ร— 94nm ร— 500nm
  • Confocal: varies by objective and zoom

2. Trace Dendritic Paths

  • Click waypoints along dendrite structures
  • Algorithm automatically finds optimal brightess paths

3. Segment Dendrites

  • Load pre-trained SAMv2 dendrite model
  • Segment individual path with SAMv2

4. Detect Spines

  • Smart multi-view detection using tube data generation
  • Angle-based matching between 2D and tubular views

5. Segment Spines

  • Fine-grained spine segmentation using specialized SAMv2 model
  • Dendrite mask overlay to suppress background signal
  • Manual point extension across frames
  • Contrasting color system for visualization

๐Ÿ”ง Model Training

Dendrite Model

cd Train-SAMv2
python train_dendrites.py --ppn 20 --pnn 10 --batch_size 32 --model_name "small"

Spine Model

python train_spines.py --model_name "small" --batch_size 16

๐Ÿ“ Data Format

Input Requirements

  • Image Format: TIFF, .d3set (to reproduce training results)
  • Dimensions: 3D volumes (Zร—Yร—X)
  • Bit Depth: 8-bit or 16-bit grayscale
  • Size: Tested up to 2048ร—2048ร—500 voxels

Output Formats

  • Paths: NumPy arrays with coordinates
  • Masks: Binary TIFF volumes

๐Ÿ“„ License

This project is licensed under the MIT License - see the LICENSE file for details.

Useful VSCode Extensions

๐Ÿ“ฌ Contact

Nipun Arora - nipun.arora@fau.de


Made with โ™ฅ๏ธ at Anki Lab ๐Ÿง โœจ

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

neuro_sam-0.1.9.tar.gz (193.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

neuro_sam-0.1.9-py3-none-any.whl (251.6 kB view details)

Uploaded Python 3

File details

Details for the file neuro_sam-0.1.9.tar.gz.

File metadata

  • Download URL: neuro_sam-0.1.9.tar.gz
  • Upload date:
  • Size: 193.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.10.19

File hashes

Hashes for neuro_sam-0.1.9.tar.gz
Algorithm Hash digest
SHA256 76eb3c45512c5f33d9117c109fe94d9cc07a148e11073a67cd2017a786e32449
MD5 aefce050ca0576dca6769c8ec3c0f8bd
BLAKE2b-256 9a8a2f6158a03012d8acf4d44ae44722d48665c170c9a88b205a64535c15f5a0

See more details on using hashes here.

File details

Details for the file neuro_sam-0.1.9-py3-none-any.whl.

File metadata

  • Download URL: neuro_sam-0.1.9-py3-none-any.whl
  • Upload date:
  • Size: 251.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.10.19

File hashes

Hashes for neuro_sam-0.1.9-py3-none-any.whl
Algorithm Hash digest
SHA256 f9f3b39186fd7c99aff5102d2cc64e2caa565ef4ef6a82aff834943e58f81d4e
MD5 5675bda50170c51bfba75c7a398e957e
BLAKE2b-256 c4f1586ef1189d81b758f9a520402d7c29b01029f1b72a44e7ca298219e1f25c

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page