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

trx_rs/formats/
trk.rs
