trx_rs/formats/tt/
mod.rs

1mod decode;
2mod mat;
3
4use std::collections::{BTreeMap, HashMap, HashSet};
5use std::path::{Path, PathBuf};
6
7use serde_json::json;
8
9use crate::dtype::DType;
10use crate::error::{Result, TrxError};
11use crate::header::Header;
12use crate::mmap_backing::vec_to_bytes;
13use crate::tractogram::Tractogram;
14use crate::trx_file::DataArray;
15
16use self::decode::{apply_affine, decode_tiny_track, TinyTrackData};
17use self::mat::read_tt_mat_records;
18
19/// Read a DSI Studio Tiny Track `.tt.gz` file into a [`Tractogram`].
20///
21/// Cluster IDs become TRX groups, and TT colors become `dpg/<group>/color.3.uint8`.
22/// Sidecar `.txt` labels are used for group names when present.
23pub fn read_tt(path: &Path) -> Result<Tractogram> {
24    if !path
25        .file_name()
26        .and_then(|name| name.to_str())
27        .is_some_and(|name| name.ends_with(".tt.gz"))
28    {
29        return Err(TrxError::Format(
30            "Tiny Track import currently supports .tt.gz inputs only".into(),
31        ));
32    }
33
34    let records = read_tt_mat_records(path)?;
35    let tt = decode_tiny_track(&records)?;
36    let sidecar_labels = read_labels_sidecar(path)?;
37    build_tractogram(tt, &sidecar_labels)
38}
39
40fn build_tractogram(tt: TinyTrackData, sidecar_labels: &[String]) -> Result<Tractogram> {
41    let mut header = Header {
42        voxel_to_rasmm: tt.affine,
43        dimensions: tt.dimensions.map(u64::from),
44        nb_streamlines: 0,
45        nb_vertices: 0,
46        extra: Default::default(),
47    };
48    if let Some(report) = tt.report {
49        header
50            .extra
51            .insert("tt_report".into(), serde_json::Value::String(report));
52    }
53    if let Some(parameter_id) = tt.parameter_id {
54        header.extra.insert(
55            "tt_parameter_id".into(),
56            serde_json::Value::String(parameter_id),
57        );
58    }
59
60    let mut tractogram = Tractogram::with_header(header);
61    let mut cluster_members: BTreeMap<u16, Vec<u32>> = BTreeMap::new();
62
63    for (index, streamline_vox) in tt.streamlines_vox.iter().enumerate() {
64        let streamline_world: Vec<[f32; 3]> = streamline_vox
65            .iter()
66            .map(|point| apply_affine(tt.affine, *point))
67            .collect();
68        tractogram.push_streamline(&streamline_world)?;
69        if let Some(cluster_id) = tt.cluster_ids.get(index).copied() {
70            cluster_members
71                .entry(cluster_id)
72                .or_default()
73                .push(index as u32);
74        }
75    }
76
77    let group_names = resolve_group_names(&cluster_members, sidecar_labels);
78    for (cluster_id, members) in &cluster_members {
79        let name = group_names
80            .get(cluster_id)
81            .ok_or_else(|| TrxError::Format("missing resolved TT group name".into()))?;
82        tractogram.insert_group(name.clone(), members.clone());
83        if let Some(&packed) = tt.colors.get(*cluster_id as usize) {
84            tractogram.insert_dpg(
85                name.clone(),
86                "color",
87                DataArray::owned_bytes(
88                    vec_to_bytes(vec![packed_color_to_rgb(packed)]),
89                    3,
90                    DType::UInt8,
91                ),
92            );
93        }
94    }
95
96    if !tt.colors.is_empty() {
97        tractogram.extra_mut().insert(
98            "tt_raw_colors".into(),
99            json!(tt
100                .colors
101                .iter()
102                .map(|color| format!("0x{color:08x}"))
103                .collect::<Vec<_>>()),
104        );
105    }
106
107    Ok(tractogram)
108}
109
110fn read_labels_sidecar(path: &Path) -> Result<Vec<String>> {
111    let mut sidecar = PathBuf::from(path);
112    let file_name = sidecar
113        .file_name()
114        .and_then(|name| name.to_str())
115        .ok_or_else(|| TrxError::Argument(format!("invalid TT path {}", path.display())))?
116        .to_string();
117
118    // Try <name>.txt first (e.g. human.tt.gz.txt), then DSI Studio convention
119    // of stripping the .tt/.tt.gz extension (e.g. human.txt).
120    let candidates = [
121        format!("{file_name}.txt"),
122        file_name
123            .strip_suffix(".tt.gz")
124            .or_else(|| file_name.strip_suffix(".tt"))
125            .map(|stem| format!("{stem}.txt"))
126            .unwrap_or_default(),
127    ];
128    for candidate in &candidates {
129        if candidate.is_empty() {
130            continue;
131        }
132        sidecar.set_file_name(candidate);
133        if sidecar.exists() {
134            let text = std::fs::read_to_string(&sidecar)?;
135            return Ok(text.lines().map(|line| line.trim().to_string()).collect());
136        }
137    }
138    Ok(Vec::new())
139}
140
141fn resolve_group_names(
142    cluster_members: &BTreeMap<u16, Vec<u32>>,
143    labels: &[String],
144) -> HashMap<u16, String> {
145    let mut used = HashSet::new();
146    let mut resolved = HashMap::new();
147    for cluster_id in cluster_members.keys() {
148        let base = labels
149            .get(*cluster_id as usize)
150            .map(|name| sanitize_group_name(name))
151            .filter(|name| !name.is_empty())
152            .unwrap_or_else(|| format!("cluster_{cluster_id}"));
153        let mut candidate = base.clone();
154        let mut suffix = 1usize;
155        while !used.insert(candidate.clone()) {
156            candidate = format!("{base}_{suffix}");
157            suffix += 1;
158        }
159        resolved.insert(*cluster_id, candidate);
160    }
161    resolved
162}
163
164fn sanitize_group_name(name: &str) -> String {
165    name.chars()
166        .map(|ch| match ch {
167            '/' | '\\' | ':' | '\0' => '_',
168            _ => ch,
169        })
170        .collect::<String>()
171        .trim()
172        .to_string()
173}
174
175fn packed_color_to_rgb(color: u32) -> [u8; 3] {
176    [
177        ((color >> 16) & 0xff) as u8,
178        ((color >> 8) & 0xff) as u8,
179        (color & 0xff) as u8,
180    ]
181}
trx_rs/formats/tt/mod.rs

trx_rs/formats/tt/
mod.rs
