trx_rs/formats/
tck.rs

1use std::fs::File;
2use std::io::{Read, Write};
3use std::path::Path;
4
5use flate2::read::MultiGzDecoder;
6use flate2::write::GzEncoder;
7use flate2::Compression;
8
9use crate::error::{Result, TrxError};
10use crate::header::Header;
11use crate::tractogram::Tractogram;
12
13const TCK_MAGIC: &str = "mrtrix tracks";
14
15/// Read an MRtrix `.tck` or `.tck.gz` file into a [`Tractogram`].
16pub fn read_tck(path: &Path, header_override: Option<Header>) -> Result<Tractogram> {
17    let bytes = read_maybe_gzip(path)?;
18    let parsed = parse_tck_bytes(&bytes)?;
19
20    let mut tractogram = Tractogram::with_header(header_override.unwrap_or(Header {
21        voxel_to_rasmm: Header::identity_affine(),
22        dimensions: [1, 1, 1],
23        nb_streamlines: 0,
24        nb_vertices: 0,
25        extra: Default::default(),
26    }));
27
28    for streamline in parsed {
29        tractogram.push_streamline(&streamline)?;
30    }
31
32    Ok(tractogram)
33}
34
35/// Write a [`Tractogram`] to `.tck` or `.tck.gz` format.
36pub fn write_tck(path: &Path, tractogram: &Tractogram) -> Result<()> {
37    let bytes = build_tck_bytes(tractogram);
38    if path
39        .file_name()
40        .and_then(|name| name.to_str())
41        .is_some_and(|name| name.ends_with(".gz"))
42    {
43        let file = File::create(path)?;
44        let mut encoder = GzEncoder::new(file, Compression::default());
45        encoder.write_all(&bytes)?;
46        encoder.finish()?;
47    } else {
48        std::fs::write(path, bytes)?;
49    }
50    Ok(())
51}
52
53fn read_maybe_gzip(path: &Path) -> Result<Vec<u8>> {
54    let mut bytes = Vec::new();
55    let mut file = File::open(path)?;
56    if path
57        .file_name()
58        .and_then(|name| name.to_str())
59        .is_some_and(|name| name.ends_with(".gz"))
60    {
61        let mut decoder = MultiGzDecoder::new(file);
62        decoder.read_to_end(&mut bytes)?;
63    } else {
64        file.read_to_end(&mut bytes)?;
65    }
66    Ok(bytes)
67}
68
69fn parse_tck_bytes(bytes: &[u8]) -> Result<Vec<Vec<[f32; 3]>>> {
70    let (data_offset, declared_count) = parse_header(bytes)?;
71    let payload = bytes
72        .get(data_offset..)
73        .ok_or_else(|| TrxError::Format("TCK file payload offset is beyond end of file".into()))?;
74
75    if payload.len() % 12 != 0 {
76        return Err(TrxError::Format(
77            "TCK payload length is not divisible by 12 bytes".into(),
78        ));
79    }
80
81    let mut streamlines = Vec::new();
82    let mut current = Vec::new();
83    let mut saw_eof = false;
84
85    for chunk in payload.chunks_exact(12) {
86        let point = [
87            f32::from_le_bytes(chunk[0..4].try_into().unwrap()),
88            f32::from_le_bytes(chunk[4..8].try_into().unwrap()),
89            f32::from_le_bytes(chunk[8..12].try_into().unwrap()),
90        ];
91
92        if point.iter().all(|value| value.is_infinite()) {
93            saw_eof = true;
94            break;
95        }
96        if point.iter().all(|value| value.is_nan()) {
97            streamlines.push(std::mem::take(&mut current));
98            continue;
99        }
100
101        current.push(point);
102    }
103
104    if !current.is_empty() {
105        streamlines.push(current);
106    }
107
108    if !saw_eof {
109        return Err(TrxError::Format("TCK file is missing EOF delimiter".into()));
110    }
111
112    if let Some(count) = declared_count {
113        let actual = u64::try_from(streamlines.len()).unwrap_or(u64::MAX);
114        if count != actual {
115            return Err(TrxError::Format(format!(
116                "TCK header count ({count}) does not match parsed streamline count ({actual})"
117            )));
118        }
119    }
120
121    Ok(streamlines)
122}
123
124fn parse_header(bytes: &[u8]) -> Result<(usize, Option<u64>)> {
125    let mut cursor = 0usize;
126    let mut lines = Vec::new();
127    while cursor < bytes.len() {
128        let line_start = cursor;
129        while cursor < bytes.len() && bytes[cursor] != b'\n' {
130            cursor += 1;
131        }
132        let line_end = cursor;
133        if cursor < bytes.len() && bytes[cursor] == b'\n' {
134            cursor += 1;
135        }
136        let line = std::str::from_utf8(&bytes[line_start..line_end])
137            .map_err(|_| TrxError::Format("TCK header is not valid UTF-8".into()))?;
138        // Some generators leave trailing spaces on header lines, and a UTF-8 BOM can
139        // appear at the beginning of the file. Accept those variants.
140        let line = line.trim_start_matches('\u{feff}').trim();
141        if line == "END" {
142            break;
143        }
144        lines.push(line.to_owned());
145    }
146
147    if lines.is_empty() || lines[0] != TCK_MAGIC {
148        return Err(TrxError::Format(
149            "file does not start with 'mrtrix tracks'".into(),
150        ));
151    }
152
153    let mut data_offset = cursor;
154    let mut declared_count = None;
155
156    for line in lines.iter().skip(1) {
157        if let Some((key, value)) = line.split_once(':') {
158            let key = key.trim();
159            let value = value.trim();
160            match key {
161                "file" => {
162                    let mut tokens = value.split_whitespace();
163                    let locator = tokens.next().unwrap_or_default();
164                    let offset = tokens
165                        .next()
166                        .ok_or_else(|| TrxError::Format("invalid TCK file offset".into()))?;
167                    if locator != "." {
168                        return Err(TrxError::Format(
169                            "only inline TCK payloads ('file: . <offset>') are supported".into(),
170                        ));
171                    }
172                    data_offset = offset.parse::<usize>().map_err(|_| {
173                        TrxError::Format(format!("invalid TCK payload offset '{offset}'"))
174                    })?;
175                }
176                "count" => {
177                    declared_count = Some(value.parse::<u64>().map_err(|_| {
178                        TrxError::Format(format!("invalid TCK streamline count '{value}'"))
179                    })?);
180                }
181                "datatype" if value != "Float32LE" => {
182                    return Err(TrxError::Format(format!(
183                        "unsupported TCK datatype '{value}', expected Float32LE"
184                    )));
185                }
186                _ => {}
187            }
188        }
189    }
190
191    Ok((data_offset, declared_count))
192}
193
194fn build_tck_bytes(tractogram: &Tractogram) -> Vec<u8> {
195    let header = build_header(tractogram.nb_streamlines());
196    let mut bytes = header.into_bytes();
197
198    for streamline in tractogram.streamlines() {
199        for point in streamline {
200            bytes.extend_from_slice(&point[0].to_le_bytes());
201            bytes.extend_from_slice(&point[1].to_le_bytes());
202            bytes.extend_from_slice(&point[2].to_le_bytes());
203        }
204        bytes.extend_from_slice(&f32::NAN.to_le_bytes());
205        bytes.extend_from_slice(&f32::NAN.to_le_bytes());
206        bytes.extend_from_slice(&f32::NAN.to_le_bytes());
207    }
208
209    bytes.extend_from_slice(&f32::INFINITY.to_le_bytes());
210    bytes.extend_from_slice(&f32::INFINITY.to_le_bytes());
211    bytes.extend_from_slice(&f32::INFINITY.to_le_bytes());
212    bytes
213}
214
215fn build_header(streamline_count: usize) -> String {
216    let mut offset = 0usize;
217    loop {
218        let header = format!(
219            "{TCK_MAGIC}\ncount: {streamline_count:010}\ndatatype: Float32LE\nfile: . {offset}\nEND\n"
220        );
221        let next_offset = header.len();
222        if next_offset == offset {
223            return header;
224        }
225        offset = next_offset;
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232
233    #[test]
234    fn header_builder_converges() {
235        let header = build_header(13);
236        let (_, count) = parse_header(header.as_bytes()).unwrap();
237        assert_eq!(count, Some(13));
238    }
239
240    #[test]
241    fn parser_accepts_padded_magic_line() {
242        let mut offset = 0usize;
243        let header = loop {
244            let header = format!(
245                "mrtrix tracks    \ncount: 0000000001\ndatatype: Float32LE\nfile: . {offset}\nEND\n"
246            );
247            let next_offset = header.len();
248            if next_offset == offset {
249                break header;
250            }
251            offset = next_offset;
252        };
253
254        let mut bytes = header.into_bytes();
255        bytes.extend_from_slice(&1.0f32.to_le_bytes());
256        bytes.extend_from_slice(&2.0f32.to_le_bytes());
257        bytes.extend_from_slice(&3.0f32.to_le_bytes());
258        bytes.extend_from_slice(&f32::NAN.to_le_bytes());
259        bytes.extend_from_slice(&f32::NAN.to_le_bytes());
260        bytes.extend_from_slice(&f32::NAN.to_le_bytes());
261        bytes.extend_from_slice(&f32::INFINITY.to_le_bytes());
262        bytes.extend_from_slice(&f32::INFINITY.to_le_bytes());
263        bytes.extend_from_slice(&f32::INFINITY.to_le_bytes());
264
265        let streamlines = parse_tck_bytes(&bytes).unwrap();
266        assert_eq!(streamlines, vec![vec![[1.0, 2.0, 3.0]]]);
267    }
268}
trx_rs/formats/tck.rs

trx_rs/formats/
tck.rs
