Skip to main content

qsm_core/segment/
labels.rs

1//! SynthSeg label tables.
2//!
3//! One table per model generation. The entries are ordered by ascending FreeSurfer label id,
4//! which is the channel order of the network's output, so index `c` of a posterior map is
5//! label [`SynthSegLabels::ids`]`[c]`.
6//!
7//! Transcribed from `data/labels_classes_priors/` in the SynthSeg distribution (Apache-2.0),
8//! with `flip` and `topology` pre-resolved from `SynthSeg/predict.py:get_flip_indices` and the
9//! topological-class file so the runtime does no table building.
10
11/// Which SynthSeg generation a set of weights belongs to.
12///
13/// The generations differ in their label set — 2.0 adds a general `csf` class (FreeSurfer id
14/// 24) — so the weights and the table have to agree.
15#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
16#[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
17pub enum SynthSegVersion {
18    /// SynthSeg 1.0 (`synthseg_1.0.h5`), 32 labels. The default because its weights ship in
19    /// the SynthSeg repository itself, so the export is reproducible from an Apache-2.0 source.
20    #[default]
21    V1,
22    /// SynthSeg 2.0 (`synthseg_2.0.h5`), 33 labels — adds `csf`. The better model, but its
23    /// weights are a separate download (FreeSurfer, or the link in the SynthSeg README).
24    V2,
25}
26
27/// The label table for one SynthSeg generation.
28pub struct SynthSegLabels {
29    /// FreeSurfer label ids, ascending; index = network output channel.
30    pub ids: &'static [i32],
31    /// Human-readable structure names, parallel to [`ids`](Self::ids).
32    pub names: &'static [&'static str],
33    /// Channel permutation applied to the left-right flipped pass before averaging: channel `c`
34    /// of the flipped posteriors corresponds to channel `flip[c]` of the unflipped ones.
35    pub flip: &'static [usize],
36    /// Topological class per channel. Channels sharing a non-zero class are cleaned up together
37    /// by keeping the largest connected component of their union; class 0 is exempt.
38    pub topology: &'static [u32],
39}
40
41impl SynthSegVersion {
42    /// The label table for this generation.
43    pub fn labels(self) -> &'static SynthSegLabels {
44        match self {
45            SynthSegVersion::V1 => &V1,
46            SynthSegVersion::V2 => &V2,
47        }
48    }
49
50    /// Number of network output channels (= number of labels, background included).
51    pub fn n_labels(self) -> usize {
52        self.labels().ids.len()
53    }
54}
55
56static V1: SynthSegLabels = SynthSegLabels {
57    ids: &[
58        0, 2, 3, 4, 5, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 26, 28, 41, 42, 43, 44, 46, 47,
59        49, 50, 51, 52, 53, 54, 58, 60,
60    ],
61    names: &[
62        "background",
63        "left cerebral white matter",
64        "left cerebral cortex",
65        "left lateral ventricle",
66        "left inferior lateral ventricle",
67        "left cerebellum white matter",
68        "left cerebellum cortex",
69        "left thalamus",
70        "left caudate",
71        "left putamen",
72        "left pallidum",
73        "3rd ventricle",
74        "4th ventricle",
75        "brain-stem",
76        "left hippocampus",
77        "left amygdala",
78        "left accumbens area",
79        "left ventral DC",
80        "right cerebral white matter",
81        "right cerebral cortex",
82        "right lateral ventricle",
83        "right inferior lateral ventricle",
84        "right cerebellum white matter",
85        "right cerebellum cortex",
86        "right thalamus",
87        "right caudate",
88        "right putamen",
89        "right pallidum",
90        "right hippocampus",
91        "right amygdala",
92        "right accumbens area",
93        "right ventral DC",
94    ],
95    flip: &[
96        0, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 11, 12, 13, 28, 29, 30, 31, 1, 2, 3, 4, 5, 6,
97        7, 8, 9, 10, 14, 15, 16, 17,
98    ],
99    topology: &[
100        0, 4, 4, 4, 4, 5, 5, 6, 7, 8, 9, 1, 2, 3, 10, 11, 12, 13, 14, 14, 14, 14, 15, 15, 16, 17,
101        18, 19, 20, 21, 22, 23,
102    ],
103};
104
105static V2: SynthSegLabels = SynthSegLabels {
106    ids: &[
107        0, 2, 3, 4, 5, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17, 18, 24, 26, 28, 41, 42, 43, 44, 46,
108        47, 49, 50, 51, 52, 53, 54, 58, 60,
109    ],
110    names: &[
111        "background",
112        "left cerebral white matter",
113        "left cerebral cortex",
114        "left lateral ventricle",
115        "left inferior lateral ventricle",
116        "left cerebellum white matter",
117        "left cerebellum cortex",
118        "left thalamus",
119        "left caudate",
120        "left putamen",
121        "left pallidum",
122        "3rd ventricle",
123        "4th ventricle",
124        "brain-stem",
125        "left hippocampus",
126        "left amygdala",
127        "csf",
128        "left accumbens area",
129        "left ventral DC",
130        "right cerebral white matter",
131        "right cerebral cortex",
132        "right lateral ventricle",
133        "right inferior lateral ventricle",
134        "right cerebellum white matter",
135        "right cerebellum cortex",
136        "right thalamus",
137        "right caudate",
138        "right putamen",
139        "right pallidum",
140        "right hippocampus",
141        "right amygdala",
142        "right accumbens area",
143        "right ventral DC",
144    ],
145    flip: &[
146        0, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 11, 12, 13, 29, 30, 16, 31, 32, 1, 2, 3, 4, 5,
147        6, 7, 8, 9, 10, 14, 15, 17, 18,
148    ],
149    topology: &[
150        0, 4, 4, 4, 4, 5, 5, 6, 7, 8, 9, 1, 2, 3, 10, 11, 0, 12, 13, 14, 14, 14, 14, 15, 15, 16,
151        17, 18, 19, 20, 21, 22, 23,
152    ],
153};
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    fn check(v: SynthSegVersion) {
160        let l = v.labels();
161        let n = l.ids.len();
162        assert_eq!(l.names.len(), n, "names length");
163        assert_eq!(l.flip.len(), n, "flip length");
164        assert_eq!(l.topology.len(), n, "topology length");
165        assert!(l.ids.windows(2).all(|w| w[0] < w[1]), "ids must be strictly ascending");
166        // Flipping twice is the identity, and background never moves.
167        assert_eq!(l.flip[0], 0);
168        for c in 0..n {
169            assert_eq!(l.flip[l.flip[c]], c, "flip is not an involution at {c}");
170        }
171        // Every left structure is paired with the matching right one.
172        for (c, name) in l.names.iter().enumerate() {
173            if let Some(rest) = name.strip_prefix("left ") {
174                assert_eq!(l.names[l.flip[c]], format!("right {rest}"));
175            }
176        }
177    }
178
179    #[test]
180    fn v1_table_is_self_consistent() {
181        check(SynthSegVersion::V1);
182        assert_eq!(SynthSegVersion::V1.n_labels(), 32);
183    }
184
185    #[test]
186    fn v2_table_is_self_consistent() {
187        check(SynthSegVersion::V2);
188        assert_eq!(SynthSegVersion::V2.n_labels(), 33);
189        // 2.0's addition over 1.0.
190        assert!(SynthSegVersion::V2.labels().ids.contains(&24));
191        assert!(!SynthSegVersion::V1.labels().ids.contains(&24));
192    }
193}