TetraN

Official PyTorch implementation of "Segmenting Fine-grained Biomedical Structures with TetraN: An N-Shaped Tetrahedral Routing Approach."

TetraN is a lightweight binary segmentation network for fine-grained biomedical structures. It combines an N-shaped encoder-decoder backbone, cross-stage global routing, Multi-block Fusers (MbFs), and a terminal Multi-scale Aggregator (MsA). The training objective applies deep supervision to the five MbF predictions and the final MsA prediction, with a tracing-aware constraint on the final output.

.
|-- assets/
|   `-- Framework.gif
|-- models/
|   |-- __init__.py
|   `-- tetran.py
|-- dataset.py
|-- losses.py
|-- train.py
|-- eval.py
|-- requirements.txt
`-- LICENSE

Python 3.10 or newer is recommended.

git clone https://github.com/BM-AI-Lab/TetraN.git
cd TetraN
pip install -r requirements.txt

Install a PyTorch build compatible with your CUDA version when the generic pip installation does not select the appropriate build. The experiments reported in the paper were run on an NVIDIA RTX 4090D GPU.

Dataset-specific preprocessing is intentionally kept outside this repository. Prepare a two-column CSV containing an image path and its binary mask path:

image,mask
images/0001.png,masks/0001.png
images/0002.png,masks/0002.png

The header is optional. Relative paths are resolved from the directory that contains the CSV file. Images are read as RGB. Every non-zero mask value is treated as foreground.

Run the paper configuration on all five folds:

python train.py \
  --data-csv data/dataset.csv \
  --output-dir outputs/dataset

The default configuration matches the implementation details in the paper:

To run one fold, use its zero-based index:

python train.py \
  --data-csv data/dataset.csv \
  --output-dir outputs/dataset \
  --fold 0

Each fold stores best.pth, last.pth, the train/validation manifests, per-epoch history, and validation metrics. cross_validation_metrics.json reports the mean and standard deviation over the executed folds.

With deep supervision enabled, TetraN returns five intermediate MbF logits and one final MsA logit. Dice+Focal loss is applied using the weights [1/16, 1/8, 1/4, 1/2, 1, 1]. Tracing Loss is applied only to the final prediction with a coefficient of 0.5, following the paper:

$$ \mathcal{L}_{base}(h_i, h_{gt}) = \mathcal{L}_{Dice}(h_i, h_{gt}) + \mathcal{L}_{Focal}(h_i, h_{gt}) $$

$$ \mathcal{L}_{total} = \lambda\cdot\mathcal{L}_{trace}(\hat{h}_{tetra}, h_{gt}) +\sum_{i=0}^{L} w_i \mathcal{L}_{base}(h_i, h_{gt}) $$

Evaluate a checkpoint on any compatible manifest:

python eval.py \
  --data-csv outputs/dataset/fold_0/val.csv \
  --checkpoint outputs/dataset/fold_0/best.pth \
  --output-dir evaluation/fold_0 \
  --profile

Add --save-predictions to write binary prediction masks. Evaluation reports Average Precision (AP), centerline Dice (clDice), Dice Similarity Coefficient (DSC), mean Intersection over Union (mIoU), and Optimal Dataset Scale (ODS). FLOPs are profiled at the selected input resolution when --profile is used.

Citation information will be added when the PRCV 2026 proceedings become available.

This project is released under the MIT License.