Skip to content

CSDLLab/NETracer

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

7 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

NETracer: A Topology-Aware Iterative Tracing Approach for Tubular Structure Extraction

ICCV 2025

Extracting tubular structures from images is a widespread and challenging task in computer vision. Existing iterative tracing methods tend to jump to adjacent branches during tracing process, leading to significant topological mistakes. The reason is that the tracing model only focuses on the estimation of discrete nodes and ignores their connection attribution.

NETracer introduces a topology-aware iterative tracing method that treats curvilinear structures as continuous node-edge sequences instead of isolated points. By jointly predicting the next node and its connecting edge, our tracer no longer "jumps" to adjacent branches even in noisy or densely packed regions.

Highlights

  • Topology-Aware Tracing: First method to seamlessly integrate connectivity features into the iterative tracing process
  • Node-Edge Estimation Network (NENet): Multi-task learning with local connectivity loss to produce future nodes and their connective edges
  • Geodesic Distance-based Search (GDS): Novel strategy to pinpoint the most suitable future node based on predicted edge cues
  • New Evaluation Metrics: Brand-new iterative tracing metrics (PE, ABL, JE) to assess local accuracy, continuity, and topological correctness
  • 2D & 3D Support: Supports both 2D images (roads, vessels) and 3D volumetric data (neurons)

Project Structure

NETracer/
├── train_2D.py                 # 2D model training and testing entry
├── train_3D.py                 # 3D model training and testing entry
├── prepare_datasets_2D.py      # 2D dataset preparation
├── prepare_datasets_3D.py      # 3D dataset preparation
├── configs/
│   ├── config_2d.py            # 2D training configuration
│   └── config_3d.py            # 3D training configuration
├── models/
│   ├── RPNet_2D.py             # 2D Point-Edge prediction network
│   └── RPNet_3D.py             # 3D Point-Edge prediction network
├── tools/
│   ├── Data_Loader_2d.py       # 2D data loader
│   ├── Data_Loader_3d.py       # 3D data loader
│   ├── Losses.py               # Loss function definitions
│   ├── dataset_2d.py           # 2D data augmentation
│   ├── dataset_3d.py           # 3D data augmentation
│   └── tracing/
│       ├── tracing_tools_2D.py # 2D iterative tracing strategy
│       └── tracing_tools_3D.py # 3D iterative tracing strategy
└── lib/
    ├── klib/                   # Basic I/O and graphics library
    ├── swclib/                 # SWC format processing library
    └── gwdt/                   # Generalized distance transform library

Installation

Requirements

  • Python >= 3.8
  • PyTorch >= 1.10
  • CUDA (recommended for GPU acceleration)

Install Dependencies

pip install torch torchvision
pip install numpy scipy scikit-image
pip install tifffile
pip install rtree
pip install GeodisTK

Method Overview

Problem Formulation

A tubular structure is a continuous curve that can be discretized into M nodes. Traditional tracing methods estimate future nodes by maximizing:

c* = argmax P(c(ti) | W[c(t1),...,c(ti-1)], I)

However, this ignores the connectivity between neighboring nodes. NETracer introduces a new optimization objective that captures the continuous edge probability:

c* = argmax ∫ P(c(t) | W[c(t1),...,c(ti-1)], I) dt

NENet Architecture

The Node-Edge Estimation Network (NENet) contains three sub-task learning blocks:

  1. Global Centerline Segmentation: Provides off-centerline stopping criteria and initial seed nodes
  2. Future Node Estimation: Predicts potential successor nodes in a probability map
  3. Local Future Edge Estimation: Provides local connectivity cues with a novel local connectivity loss

The overall loss function:

L = L_cl + L_node + L_edge
L_edge = L_e + λ * L_conn

Geodesic Distance-based Search (GDS)

When multiple potential future nodes exist, GDS strategy:

  1. Extracts all potential future nodes from the probability map
  2. Starts from current node and traces along predicted edge using geodesic distance
  3. Finds the most suitable future node with topological reliability

Supported Datasets

Type Dataset Description Reference
2D Massachusetts Roads Aerial road images Mnih 2013
2D DRIVE Retinal vessel images Staal et al. 2004
2D CHASE_DB1 Retinal vessel images Carballal et al. 2018
3D fMOST-VTA Brain neuron imaging Li et al. 2010
3D DIADEM-CCF Brain neuron imaging Brown et al. 2011

Usage

1. Prepare Dataset

First, prepare training data by processing raw SWC files and images:

For 2D datasets:

python prepare_datasets_2D.py

For 3D datasets:

python prepare_datasets_3D.py

This will generate paired input images and multiple label maps including:

  • node_img.tif - Input image patch
  • node_skl.tif - Centerline label
  • node_walk.tif - Walked path mask
  • node_point.tif - Next node position label
  • node_edge.tif - Edge connection label

2. Configure Training

Edit the configuration files in configs/ directory:

config_2d.py key parameters:

batch_size = 256        # Batch size for training
lr = 3e-4               # Learning rate
nb_epoch = 30           # Number of training epochs
num_vec = 50            # Vector quantization bins
patch_size = 64         # Input patch size (64x64)
r = 2                   # Seed radius

config_3d.py key parameters:

batch_size = 32         # Batch size for training
patch_size = [16, 64, 64]  # Input patch size (D x H x W)
stride = (8, 48, 48)    # Prediction stride

3. Train Model

For 2D model:

python train_2D.py --mode train

For 3D model:

python train_3D.py --mode train

4. Test Model

For 2D model:

python train_2D.py --mode test

For 3D model:

# First generate centerline and boundary distance maps
python train_3D.py --mode test_dis

# Then perform iterative tracing
python train_3D.py --mode test

Evaluation Metrics

Global Coverage Metrics

  • SSD-F1: Measures the number of nodes matched between test and gold branches
  • L-F1 (Length-F1): Measures the number of matched edges

Iterative Tracing Metrics (Proposed)

  • PE (Position Error): Euclidean distances between matching nodes, evaluates local accuracy
  • ABL (Average Branch Length): Evaluates continuity, higher means longer and more complete branches
  • JE (Jump Error): Counts critical topological errors when tracing jumps to wrong branches

Results

NETracer outperforms existing segmentation and tracing methods on all datasets:

2D Datasets

Dataset Method SSD-F1 L-F1 PE ↓ JE ↓ ABL ↑
ROAD NETracer 0.718 0.836 1.263 8.2 234.4
DRIVE NETracer 0.925 0.824 0.633 1.9 74.4
CHASE_DB1 NETracer 0.837 0.860 0.771 4.3 131.6

3D Datasets

Dataset Method SSD-F1 L-F1 PE ↓ JE ↓ ABL ↑
fMOST-VTA NETracer 0.919 0.913 0.693 2.3 219.3
DIADEM-CCF NETracer 0.792 0.755 2.491 5.1 107.1

Model Architecture

The NENet (Node-Edge Estimation Network) uses a V-Net backbone with multi-head outputs:

  • Input: Raw image (64×64 for 2D, 32×64×64 for 3D) + historical walked path mask
  • Backbone: V-Net architecture (lightweight and robust)
  • Outputs:
    • Global centerline segmentation (with weighted Cross-Entropy loss, β=0.9)
    • Future node probability map (predicts nodes within radius r)
    • Future edge probability map (with local connectivity loss)

Training Pipeline

  1. Data Preparation: Generate training samples with path correction (slight noise injection for robustness)
  2. Multi-Task Learning: Joint optimization of centerline, node, and edge prediction
  3. Loss Functions:
    • Centerline: Weighted Cross-Entropy (CE) loss
    • Node: Weighted CE loss
    • Edge: CE loss + Local connectivity loss (L_conn)
  4. Optimization: Adam optimizer with cosine annealing LR schedule

Inference Pipeline

  1. Seed Initialization: Randomly select seed nodes on predicted centerline (threshold=0.5)
  2. Iterative Tracing:
    • Extract patch around current node
    • Network predicts future node probability map and edge probability map
    • Extract potential future node set Q from probability map
    • Apply GDS strategy: trace along predicted edge using geodesic distance
    • Find most suitable future node from Q with topological reliability
    • Add new node to SWC tree structure
    • Repeat until off-centerline stopping criteria met
  3. Post-processing:
    • Remove isolated noise branches
    • Adjust node coordinates to original image space
    • Save final SWC file

Output Format

The tracing results are saved in SWC format, which is the standard format for neuronal morphology:

# SWC format: n T x y z R P
# n: node ID
# T: node type
# x, y, z: coordinates
# R: radius
# P: parent node ID (-1 for root)

Environment

  • Hardware: 64 GB memory, Intel Xeon Platinum 8383C CPU, 2× NVIDIA GeForce RTX3090 GPUs
  • Software: PyTorch 1.13.0, Python 3.8+
  • Training Time: ~8 hours

Citation

If you found this project useful, you can kindly give us a star ⭐, or cite us in your work 📖:

We would also like to thank Dr. Ting Zhao for his contributions to this project.

@inproceedings{liu2025netracer,
  title={NETracer: A Topology-Aware Iterative Tracing Approach for Tubular Structure Extraction},
  author={Liu, Chao and Jiang, Yangbo and Zheng, Nenggan},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  pages={20593--20602},
  year={2025}
}

Acknowledgments

  • V3D - 3D visualization and analysis
  • PyNeval - Python toolbox for neuron reconstruction evaluation

License

This project is released under the MIT License.

About

NETracer: A Topology-Aware Iterative Tracing Approach for Tubular Structure Extraction

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors