gif-python/run_registration_parcellation_basic.py at main · cai4cai/gif-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
"""
Basic script to run multi-atlas segmentation on a single subject.
"""
import os
from glob import glob

from src.multi_atlas.multi_atlas_segmentation import multi_atlas_segmentation

# ----------------------
# Configuration
# ----------------------
# Input: single subject directory
SUBJECT_DIR = "data/input/subject-01"

# Atlases: Mindboggle atlas directory
ATLAS_ROOT = "data/atlases/Mindboggle101"

# Output: where to save results
RESULTS_ROOT = "results"

# ----------------------
# Run segmentation
# ----------------------
def run_segmentation():
    """Run multi-atlas segmentation on a single subject."""
    subject_name = os.path.basename(SUBJECT_DIR)
    print(f"Processing: {subject_name}")

    # Input file
    img_path = os.path.join(SUBJECT_DIR, "t1weighted_brain.nii.gz")

    # Output directory
    results_dir = os.path.join(RESULTS_ROOT, subject_name)
    os.makedirs(results_dir, exist_ok=True)

    # Check if already processed
    parcellation_file = os.path.join(results_dir, "final_parcellation.nii.gz")
    if os.path.isfile(parcellation_file):
        print("Already processed - skipping")
        return

    # Collect atlas paths
    atlas_dirs = [d for d in glob(os.path.join(ATLAS_ROOT, "*")) if os.path.isdir(d)]
    atlas_paths_dicts_list = [
        {
            'name': os.path.basename(atlas_dir),
            'img_path': os.path.join(atlas_dir, 't1weighted.nii.gz'),
            'seg_path': os.path.join(atlas_dir, 'labels.DKT31.manual+aseg_cleaned.nii.gz')
        }
        for atlas_dir in atlas_dirs
    ]

    # Structure and tissue info CSVs
    structure_info_csv = os.path.join(ATLAS_ROOT, "structures_info.csv")
    tissue_info_csv = os.path.join(ATLAS_ROOT, "tissues_info.csv")

    # Run segmentation
    print(f"Running segmentation with {len(atlas_paths_dicts_list)} atlases...")

    _ = multi_atlas_segmentation(
        img_path=img_path,
        mask_path=None,
        atlas_paths_dicts_list=atlas_paths_dicts_list,
        structure_info_csv_path=structure_info_csv,
        tissue_info_csv_path=tissue_info_csv,
        save_dir=results_dir
    )

    print(f"Completed! Results saved to: {results_dir}")

# ----------------------
# Main
# ----------------------
if __name__ == "__main__":
    print("="*60)
    print(f"Subject: {SUBJECT_DIR}")
    print(f"Atlases: {ATLAS_ROOT}")
    print(f"Output: {RESULTS_ROOT}")
    print("="*60)

    run_segmentation()

    print("="*60)
    print("Done!")
    print("="*60)