trx_rs/formats/
trk.rs

1use std::collections::HashMap;
2use std::fs::File;
3use std::io::Read;
4use std::path::Path;
5
6use flate2::read::MultiGzDecoder;
7
8use crate::any_trx_file::AnyTrxFile;
9use crate::dtype::DType;
10use crate::error::{Result, TrxError};
11use crate::header::Header;
12use crate::mmap_backing::{vec_to_bytes, MmapBacking};
13use crate::tractogram::Tractogram;
14use crate::trx_file::{DataArray, TrxFile, TrxParts};
15
16const TRK_HEADER_SIZE: usize = 1000;
17const MAX_NAMED_SCALARS_PER_POINT: usize = 10;
18const MAX_NAMED_PROPERTIES_PER_STREAMLINE: usize = 10;
19
20/// Read a TrackVis `.trk` file into a [`Tractogram`].
21///
22/// Returns an error if the file contains TrackVis scalars or properties
23/// (use [`convert_trk_to_trx`] to preserve those).
24pub fn read_trk(path: &Path, _header_override: Option<Header>) -> Result<Tractogram> {
25    let parsed = parse_trk(path)?;
26
27    let mut tractogram = Tractogram::with_header(parsed.header);
28    for streamline in parsed.streamlines {
29        tractogram.push_streamline(&streamline)?;
30    }
31    for (name, array) in parsed.dpv {
32        tractogram.insert_dpv(name, array);
33    }
34    for (name, array) in parsed.dps {
35        tractogram.insert_dps(name, array);
36    }
37    Ok(tractogram)
38}
39
40/// Convert a TrackVis `.trk` file directly to TRX, preserving scalars and properties.
41pub fn convert_trk_to_trx(
42    input: &Path,
43    output: &Path,
44    options: &crate::formats::ConversionOptions,
45) -> Result<()> {
46    let parsed = parse_trk(input)?;
47    let positions: Vec<[f32; 3]> = parsed.streamlines.iter().flatten().copied().collect();
48    let offsets = build_offsets(&parsed.streamlines)?;
49
50    let file = TrxFile::from_parts(TrxParts {
51        header: Header {
52            nb_streamlines: parsed.streamlines.len() as u64,
53            nb_vertices: positions.len() as u64,
54            ..parsed.header
55        },
56        positions_backing: MmapBacking::Owned(vec_to_bytes(positions)),
57        offsets_backing: MmapBacking::Owned(vec_to_bytes(offsets)),
58        dps: parsed.dps,
59        dpv: parsed.dpv,
60        groups: HashMap::new(),
61        dpg: HashMap::new(),
62        tempdir: None,
63    });
64
65    let any = AnyTrxFile::F32(file);
66    let any = if options.trx_positions_dtype == DType::Float32 {
67        any
68    } else {
69        any.convert_positions_dtype(options.trx_positions_dtype)?
70    };
71    any.save(output)
72}
73
74struct ParsedTrk {
75    header: Header,
76    streamlines: Vec<Vec<[f32; 3]>>,
77    dpv: HashMap<String, DataArray>,
78    dps: HashMap<String, DataArray>,
79}
80
81#[derive(Clone, Copy)]
82enum Endianness {
83    Little,
84    Big,
85}
86
87impl Endianness {
88    fn read_i16(self, bytes: &[u8]) -> i16 {
89        match self {
90            Endianness::Little => i16::from_le_bytes(bytes.try_into().unwrap()),
91            Endianness::Big => i16::from_be_bytes(bytes.try_into().unwrap()),
92        }
93    }
94
95    fn read_i32(self, bytes: &[u8]) -> i32 {
96        match self {
97            Endianness::Little => i32::from_le_bytes(bytes.try_into().unwrap()),
98            Endianness::Big => i32::from_be_bytes(bytes.try_into().unwrap()),
99        }
100    }
101
102    fn read_f32(self, bytes: &[u8]) -> f32 {
103        match self {
104            Endianness::Little => f32::from_le_bytes(bytes.try_into().unwrap()),
105            Endianness::Big => f32::from_be_bytes(bytes.try_into().unwrap()),
106        }
107    }
108}
109
110#[derive(Clone)]
111struct SliceSpec {
112    name: String,
113    start: usize,
114    len: usize,
115}
116
117fn parse_trk(path: &Path) -> Result<ParsedTrk> {
118    let bytes = read_maybe_gzip(path)?;
119    if bytes.len() < TRK_HEADER_SIZE {
120        return Err(TrxError::Format(
121            "TRK file is smaller than the 1000-byte header".into(),
122        ));
123    }
124    if &bytes[..5] != b"TRACK" {
125        return Err(TrxError::Format(
126            "file does not start with TrackVis magic".into(),
127        ));
128    }
129
130    let header_bytes = &bytes[..TRK_HEADER_SIZE];
131    let endian = detect_endianness(header_bytes)?;
132    let version = endian.read_i32(&header_bytes[992..996]);
133    if version != 2 {
134        return Err(TrxError::Format(format!(
135            "unsupported TrackVis version {version}; only v2 is supported"
136        )));
137    }
138
139    let dimensions = [
140        parse_positive_i16(endian.read_i16(&header_bytes[6..8]), "dim[0]")? as usize,
141        parse_positive_i16(endian.read_i16(&header_bytes[8..10]), "dim[1]")? as usize,
142        parse_positive_i16(endian.read_i16(&header_bytes[10..12]), "dim[2]")? as usize,
143    ];
144    let voxel_sizes = [
145        parse_positive_f32(endian.read_f32(&header_bytes[12..16]), "voxel_size[0]")?,
146        parse_positive_f32(endian.read_f32(&header_bytes[16..20]), "voxel_size[1]")?,
147        parse_positive_f32(endian.read_f32(&header_bytes[20..24]), "voxel_size[2]")?,
148    ];
149
150    let n_scalars = parse_nonnegative_i16(endian.read_i16(&header_bytes[36..38]), "n_scalars")?;
151    let n_properties =
152        parse_nonnegative_i16(endian.read_i16(&header_bytes[238..240]), "n_properties")?;
153
154    let scalar_specs = parse_name_specs(
155        &header_bytes[38..238],
156        n_scalars,
157        MAX_NAMED_SCALARS_PER_POINT,
158        "scalars",
159    )?;
160    let property_specs = parse_name_specs(
161        &header_bytes[240..440],
162        n_properties,
163        MAX_NAMED_PROPERTIES_PER_STREAMLINE,
164        "properties",
165    )?;
166
167    let voxel_to_rasmm_f32 = parse_affine_f32(endian, &header_bytes[440..504])?;
168    if affine_is_all_zero(&voxel_to_rasmm_f32) {
169        return Err(TrxError::Format(
170            "TRK vox_to_ras is missing or zero; convert it with a more permissive tool first"
171                .into(),
172        ));
173    }
174    let affine_codes = affine_to_axcodes(&voxel_to_rasmm_f32)?;
175    let header_codes = parse_voxel_order(&header_bytes[948..952])?;
176
177    let declared_count = endian.read_i32(&header_bytes[988..992]);
178    let mut cursor = TRK_HEADER_SIZE;
179    let mut count = 0usize;
180    let mut streamlines = Vec::new();
181    let mut dpv_buffers = allocate_field_buffers(&scalar_specs);
182    let mut dps_buffers = allocate_field_buffers(&property_specs);
183
184    while cursor < bytes.len() {
185        if declared_count > 0 && count >= declared_count as usize {
186            break;
187        }
188
189        if cursor + 4 > bytes.len() {
190            return Err(TrxError::Format(
191                "TRK payload ended while reading streamline length".into(),
192            ));
193        }
194        let len = endian.read_i32(&bytes[cursor..cursor + 4]);
195        cursor += 4;
196        let len = usize::try_from(len)
197            .map_err(|_| TrxError::Format("TRK streamline length cannot be negative".into()))?;
198
199        let mut streamline = Vec::with_capacity(len);
200        for _ in 0..len {
201            if cursor + 12 > bytes.len() {
202                return Err(TrxError::Format(
203                    "TRK payload ended while reading streamline points".into(),
204                ));
205            }
206            let point_voxmm = [
207                endian.read_f32(&bytes[cursor..cursor + 4]),
208                endian.read_f32(&bytes[cursor + 4..cursor + 8]),
209                endian.read_f32(&bytes[cursor + 8..cursor + 12]),
210            ];
211            cursor += 12;
212            let point_world = trackvis_to_rasmm(
213                point_voxmm,
214                voxel_sizes,
215                dimensions,
216                header_codes,
217                affine_codes,
218                &voxel_to_rasmm_f32,
219            )?;
220            streamline.push(point_world);
221
222            if n_scalars > 0 {
223                let scalar_values =
224                    read_f32_row(&bytes, &mut cursor, n_scalars, endian, "scalars")?;
225                append_slices(&mut dpv_buffers, &scalar_specs, &scalar_values);
226            }
227        }
228
229        if n_properties > 0 {
230            let property_values =
231                read_f32_row(&bytes, &mut cursor, n_properties, endian, "properties")?;
232            append_slices(&mut dps_buffers, &property_specs, &property_values);
233        }
234
235        streamlines.push(streamline);
236        count += 1;
237    }
238
239    if declared_count > 0 && count != declared_count as usize {
240        return Err(TrxError::Format(format!(
241            "TRK header declares {declared_count} streamlines but parsed {count}"
242        )));
243    }
244
245    let header = Header {
246        voxel_to_rasmm: voxel_to_rasmm_f32.map(|row| row.map(f64::from)),
247        dimensions: dimensions.map(|value| value as u64),
248        nb_streamlines: streamlines.len() as u64,
249        nb_vertices: streamlines.iter().map(Vec::len).sum::<usize>() as u64,
250        extra: Default::default(),
251    };
252
253    Ok(ParsedTrk {
254        header,
255        streamlines,
256        dpv: finalize_field_buffers(dpv_buffers),
257        dps: finalize_field_buffers(dps_buffers),
258    })
259}
260
261fn detect_endianness(header: &[u8]) -> Result<Endianness> {
262    let le = i32::from_le_bytes(header[996..1000].try_into().unwrap());
263    if le == TRK_HEADER_SIZE as i32 {
264        return Ok(Endianness::Little);
265    }
266    let be = i32::from_be_bytes(header[996..1000].try_into().unwrap());
267    if be == TRK_HEADER_SIZE as i32 {
268        return Ok(Endianness::Big);
269    }
270    Err(TrxError::Format("TRK header size is invalid".into()))
271}
272
273fn parse_positive_i16(value: i16, label: &str) -> Result<i16> {
274    if value <= 0 {
275        return Err(TrxError::Format(format!(
276            "TRK {label} must be positive, got {value}"
277        )));
278    }
279    Ok(value)
280}
281
282fn parse_nonnegative_i16(value: i16, label: &str) -> Result<usize> {
283    usize::try_from(value)
284        .map_err(|_| TrxError::Format(format!("TRK {label} cannot be negative, got {value}")))
285}
286
287fn parse_positive_f32(value: f32, label: &str) -> Result<f32> {
288    if !value.is_finite() || value <= 0.0 {
289        return Err(TrxError::Format(format!(
290            "TRK {label} must be positive, got {value}"
291        )));
292    }
293    Ok(value)
294}
295
296fn parse_affine_f32(endian: Endianness, bytes: &[u8]) -> Result<[[f32; 4]; 4]> {
297    let mut affine = [[0.0f32; 4]; 4];
298    for (row, row_values) in affine.iter_mut().enumerate() {
299        for (col, value) in row_values.iter_mut().enumerate() {
300            let start = (row * 4 + col) * 4;
301            *value = endian.read_f32(&bytes[start..start + 4]);
302        }
303    }
304    if !affine.iter().flatten().all(|value| value.is_finite()) {
305        return Err(TrxError::Format(
306            "TRK vox_to_ras contains non-finite values".into(),
307        ));
308    }
309    Ok(affine)
310}
311
312fn affine_is_all_zero(affine: &[[f32; 4]; 4]) -> bool {
313    affine
314        .iter()
315        .flatten()
316        .all(|value| value.abs() <= f32::EPSILON)
317}
318
319fn parse_voxel_order(bytes: &[u8]) -> Result<[char; 3]> {
320    let text = bytes
321        .iter()
322        .copied()
323        .take_while(|byte| *byte != 0)
324        .map(char::from)
325        .collect::<String>()
326        .trim()
327        .to_ascii_uppercase();
328    let chars: Vec<char> = text.chars().collect();
329    if chars.len() < 3 {
330        return Err(TrxError::Format(
331            "TRK voxel order is missing or incomplete; convert it with another tool first".into(),
332        ));
333    }
334    Ok([chars[0], chars[1], chars[2]])
335}
336
337fn affine_to_axcodes(affine: &[[f32; 4]; 4]) -> Result<[char; 3]> {
338    let mut used_world_axes = [false; 3];
339    let mut out = ['R'; 3];
340    for (voxel_axis, slot) in out.iter_mut().enumerate() {
341        let mut best_axis = 0usize;
342        let mut best_value = 0.0f32;
343        for (world_axis, row) in affine.iter().enumerate().take(3) {
344            let value = row[voxel_axis].abs();
345            if value > best_value {
346                best_value = value;
347                best_axis = world_axis;
348            }
349        }
350        if best_value <= f32::EPSILON || used_world_axes[best_axis] {
351            return Err(TrxError::Format(
352                "TRK vox_to_ras affine is ambiguous or singular".into(),
353            ));
354        }
355        used_world_axes[best_axis] = true;
356        let sign = affine[best_axis][voxel_axis].signum();
357        *slot = match (best_axis, sign >= 0.0) {
358            (0, true) => 'R',
359            (0, false) => 'L',
360            (1, true) => 'A',
361            (1, false) => 'P',
362            (2, true) => 'S',
363            (2, false) => 'I',
364            _ => unreachable!(),
365        };
366    }
367    Ok(out)
368}
369
370fn parse_name_specs(
371    bytes: &[u8],
372    total: usize,
373    max_entries: usize,
374    fallback_name: &str,
375) -> Result<Vec<SliceSpec>> {
376    let mut specs = Vec::new();
377    let mut cursor = 0usize;
378    for entry in 0..max_entries {
379        let start = entry * 20;
380        let end = start + 20;
381        let (name, count) = decode_name_field(&bytes[start..end])?;
382        if count == 0 {
383            continue;
384        }
385        specs.push(SliceSpec {
386            name,
387            start: cursor,
388            len: count,
389        });
390        cursor += count;
391    }
392
393    if cursor < total {
394        specs.push(SliceSpec {
395            name: fallback_name.to_string(),
396            start: cursor,
397            len: total - cursor,
398        });
399        cursor = total;
400    }
401    if cursor != total {
402        return Err(TrxError::Format(format!(
403            "TrackVis named field layout is inconsistent with declared column count {total}"
404        )));
405    }
406    Ok(specs)
407}
408
409fn decode_name_field(bytes: &[u8]) -> Result<(String, usize)> {
410    let decoded = bytes.iter().map(|&byte| byte as char).collect::<String>();
411    let trimmed = decoded.trim_end_matches('\0');
412    if trimmed.is_empty() {
413        return Ok((String::new(), 0));
414    }
415
416    let mut parts = trimmed.split('\0');
417    let name = parts.next().unwrap().to_string();
418    let count = match parts.next() {
419        Some(count) => count
420            .parse::<usize>()
421            .map_err(|_| TrxError::Format(format!("invalid TrackVis name encoding '{trimmed}'")))?,
422        None => 1,
423    };
424    if parts.next().is_some() {
425        return Err(TrxError::Format(format!(
426            "invalid TrackVis name encoding '{trimmed}'"
427        )));
428    }
429    Ok((name, count))
430}
431
432fn read_f32_row(
433    bytes: &[u8],
434    cursor: &mut usize,
435    count: usize,
436    endian: Endianness,
437    label: &str,
438) -> Result<Vec<f32>> {
439    let byte_count = count
440        .checked_mul(4)
441        .ok_or_else(|| TrxError::Format(format!("TRK {label} row is too large")))?;
442    if *cursor + byte_count > bytes.len() {
443        return Err(TrxError::Format(format!(
444            "TRK payload ended while reading {label}"
445        )));
446    }
447    let row = (0..count)
448        .map(|index| {
449            let start = *cursor + index * 4;
450            endian.read_f32(&bytes[start..start + 4])
451        })
452        .collect();
453    *cursor += byte_count;
454    Ok(row)
455}
456
457fn allocate_field_buffers(specs: &[SliceSpec]) -> HashMap<String, (usize, Vec<f32>)> {
458    specs
459        .iter()
460        .map(|spec| (spec.name.clone(), (spec.len, Vec::new())))
461        .collect()
462}
463
464fn append_slices(
465    buffers: &mut HashMap<String, (usize, Vec<f32>)>,
466    specs: &[SliceSpec],
467    row: &[f32],
468) {
469    for spec in specs {
470        if let Some((_, values)) = buffers.get_mut(&spec.name) {
471            values.extend_from_slice(&row[spec.start..spec.start + spec.len]);
472        }
473    }
474}
475
476fn finalize_field_buffers(
477    buffers: HashMap<String, (usize, Vec<f32>)>,
478) -> HashMap<String, DataArray> {
479    buffers
480        .into_iter()
481        .map(|(name, (ncols, values))| {
482            (
483                name,
484                DataArray::owned_bytes(vec_to_bytes(values), ncols, DType::Float32),
485            )
486        })
487        .collect()
488}
489
490fn trackvis_to_rasmm(
491    point_voxmm: [f32; 3],
492    voxel_sizes: [f32; 3],
493    dimensions: [usize; 3],
494    header_codes: [char; 3],
495    affine_codes: [char; 3],
496    voxel_to_rasmm: &[[f32; 4]; 4],
497) -> Result<[f32; 3]> {
498    let voxel_center = [
499        point_voxmm[0] / voxel_sizes[0] - 0.5,
500        point_voxmm[1] / voxel_sizes[1] - 0.5,
501        point_voxmm[2] / voxel_sizes[2] - 0.5,
502    ];
503    let oriented = reorient_voxel_coords(voxel_center, dimensions, header_codes, affine_codes)?;
504    Ok(apply_affine(voxel_to_rasmm, oriented))
505}
506
507fn reorient_voxel_coords(
508    point: [f32; 3],
509    dimensions: [usize; 3],
510    from: [char; 3],
511    to: [char; 3],
512) -> Result<[f32; 3]> {
513    let mut out = [0.0f32; 3];
514    for (dst_axis, dst_code) in to.iter().enumerate() {
515        let dst_family = axis_family(*dst_code)?;
516        let dst_sign = axis_sign(*dst_code)?;
517        let src_axis = from
518            .iter()
519            .position(|code| axis_family(*code).ok() == Some(dst_family))
520            .ok_or_else(|| {
521                TrxError::Format("TRK voxel order does not match affine orientation".into())
522            })?;
523        let src_sign = axis_sign(from[src_axis])?;
524        out[dst_axis] = if src_sign == dst_sign {
525            point[src_axis]
526        } else {
527            (dimensions[src_axis] as f32 - 1.0) - point[src_axis]
528        };
529    }
530    Ok(out)
531}
532
533fn axis_family(code: char) -> Result<usize> {
534    match code {
535        'L' | 'R' => Ok(0),
536        'P' | 'A' => Ok(1),
537        'I' | 'S' => Ok(2),
538        other => Err(TrxError::Format(format!(
539            "unsupported anatomical axis code '{other}'"
540        ))),
541    }
542}
543
544fn axis_sign(code: char) -> Result<i8> {
545    match code {
546        'R' | 'A' | 'S' => Ok(1),
547        'L' | 'P' | 'I' => Ok(-1),
548        other => Err(TrxError::Format(format!(
549            "unsupported anatomical axis code '{other}'"
550        ))),
551    }
552}
553
554fn apply_affine(affine: &[[f32; 4]; 4], point: [f32; 3]) -> [f32; 3] {
555    [
556        affine[0][0] * point[0] + affine[0][1] * point[1] + affine[0][2] * point[2] + affine[0][3],
557        affine[1][0] * point[0] + affine[1][1] * point[1] + affine[1][2] * point[2] + affine[1][3],
558        affine[2][0] * point[0] + affine[2][1] * point[1] + affine[2][2] * point[2] + affine[2][3],
559    ]
560}
561
562fn build_offsets(streamlines: &[Vec<[f32; 3]>]) -> Result<Vec<u32>> {
563    let mut offsets = Vec::with_capacity(streamlines.len() + 1);
564    offsets.push(0);
565    let mut total = 0usize;
566    for streamline in streamlines {
567        total += streamline.len();
568        offsets
569            .push(u32::try_from(total).map_err(|_| {
570                TrxError::Format("TRK file has more than u32::MAX vertices".into())
571            })?);
572    }
573    Ok(offsets)
574}
575
576fn read_maybe_gzip(path: &Path) -> Result<Vec<u8>> {
577    let mut bytes = Vec::new();
578    let file = File::open(path)?;
579    if path
580        .file_name()
581        .and_then(|name| name.to_str())
582        .is_some_and(|name| name.ends_with(".gz"))
583    {
584        let mut decoder = MultiGzDecoder::new(file);
585        decoder.read_to_end(&mut bytes)?;
586    } else {
587        let mut file = file;
588        file.read_to_end(&mut bytes)?;
589    }
590    Ok(bytes)
591}
trx_rs/formats/trk.rs

trx_rs/formats/
trk.rs
