Segmentation of Non-Human Primate Cerebrovascular Images from Synchrotron Radiation Micro-Tomography Using Transfer Learning and Attention U-Net

This repository implements a binary segmentation framework for synchrotron radiation (SR) cerebrovascular data of non-human primates, utilizing Transfer Learning and an Attention U-Net architecture.Data：Pre-training Data: DeepVessel (Large-scale synthetic data)Fine-tuning Data: Synchrotron radiation micro-CT imaging data of macaque cerebrovasculature.Methodology:Binary segmentation is performed on Z-axis slices of the 3D volumetric data.1. Pre-training PhaseThe model is fully trained using the large-scale synthetic dataset, DeepVessel.Loss Function: Sum of Binary Cross Entropy (BCE) Loss and Dice Loss.2. Fine-tuning PhaseNormalization: Global percentile normalization is applied to the Synchrotron Radiation data.Training Strategy: The encoder is frozen. The model is fine-tuned on the new dataset using a specific learning rate.Weighting Strategy: A hierarchical compound weighting strategy is designed to be sensitive to small targets and boundaries, generating a specific weight map.Loss Function: Weighted BCE Loss (using the generated weight map) + Dice Loss.Project Structure: AttentionUNet：├── configs│  └── config.py           # Configuration parameters├── dataset              # Data loading and processing├── debug_final_check         # Validation of raw data, labels, and weight maps for fine-tuning├── models              # Model architectures│  ├── attention_unet.py│  └── blocks.py├── normalization_plots_nifti_svg   # Visualization of normalization schemes├── results│  └── 20251120│    ├── finetune_2        # Fine-tuned model weights and prediction results│    ├── best_model.pth      # Pre-trained model weights│    └── training_history.png   # Pre-training loss curves├── utils               # Utilities (Loss functions, etc.)├── compute_states.py         # Calculates p1 and p99 for global percentile normalization├── count_efficiency.py        # Calculates model efficiency (FLOPs/Params)├── evaluate_metrics.py        # Quantitative evaluation of prediction results├── finetune.py            # Main fine-tuning script├── metrics_finetune_final_v5.csv   # Final prediction metrics├── metrics_finetune_final_v5_fold0.csv # Single-fold metrics (used for ablation studies)├── metrics_nnunet_2d_final.csv    # Baseline metrics (nnU-Net 2D)├── metrics_nnunet_3d_final.csv    # Baseline metrics (nnU-Net 3D)├── plot_3d_results.py        # 3D visualization of segmentation results├── plot_error_map.py         # Cross-sectional error map visualization├── plot_normalization.py       # Script to plot normalization schematics├── predict.py            # Inference script├── predict_fold0.py         # Single-fold inference script└── run_all_folds.sh         # Shell script for 5-fold cross-validation fine-tuningAttentionUNet_ablation_01: Ablation study on the Transfer Learning strategy.AttentionUNet_ablation_02: Ablation study on the small target and boundary-aware weighting strategy.nnUNet：├── Data│  ├── nnUNet_preprocessed      # Preprocessed data used by nnU-Net│  ├── nnUNet_raw│  │  └── Dataset1120_MMSBrainBin8│  │    ├── imagesTr       # Training images (Raw)│  │    ├── imagesTs       # Testing images (Raw)│  │    ├── labelsTr       # Training labels│  │    ├── labelsTs       # Testing labels│  │    └── dataset.json     # Dataset descriptor file│  ├── nnUNet_results        # 5-fold training results│  └── predictions          # Inference results├── error_map             # Cross-sectional error map visualization├── vis_3d_components_nnunet2d    # 3D visualization of nnU-Net 2D results├── vis_3d_components_nnunet3d    # 3D visualization of nnU-Net 3D results├── 1118_MMSBrainBin8_data_analysis.csv # Descriptive statistics for all datasets├── evaluate_metrics.py        # Quantitative evaluation script├── plot_3d_results.py        # 3D visualization script├── prepare_dataset.py        # Script to calculate dataset descriptive statistics├── rename_dataset.py         # Renames dataset to nnU-Net compatible format└── test_nnunet_efficiency.py     # Calculates nnU-Net model efficiency