trx_rs/formats/
vtk.rs

1use std::path::Path;
2
3use crate::error::{Result, TrxError};
4use crate::header::Header;
5use crate::tractogram::Tractogram;
6
7/// Controls how VTK coordinate spaces are interpreted during import.
8#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
9pub enum VtkCoordinateMode {
10    /// Read the VTK header comment for `SPACE=`, warn if absent.
11    HeaderOrWarn,
12    /// Assume coordinates are in RAS space (the default).
13    #[default]
14    AssumeRas,
15    /// Assume coordinates are in LPS space.
16    AssumeLps,
17}
18
19/// The anatomical coordinate space declared or inferred for VTK data.
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21pub enum VtkCoordinateSpace {
22    /// Right-Anterior-Superior.
23    Ras,
24    /// Left-Posterior-Superior.
25    Lps,
26}
27
28pub fn read_vtk(
29    path: &Path,
30    header_override: Option<Header>,
31    coordinate_mode: VtkCoordinateMode,
32) -> Result<Tractogram> {
33    let bytes = std::fs::read(path)?;
34    let (positions, offsets) = parse_vtk_bytes(&bytes, coordinate_mode)?;
35
36    let header = header_override.unwrap_or(Header {
37        voxel_to_rasmm: Header::identity_affine(),
38        dimensions: [1, 1, 1],
39        nb_streamlines: 0,
40        nb_vertices: 0,
41        extra: Default::default(),
42    });
43
44    Ok(Tractogram::from_positions_and_offsets(
45        positions, offsets, header,
46    ))
47}
48
49use std::io::Write;
50
51pub fn write_vtk(path: &Path, tractogram: &Tractogram) -> Result<()> {
52    let file = std::fs::File::create(path)?;
53    let mut file = std::io::BufWriter::new(file);
54
55    let mut header = String::new();
56    header.push_str("# vtk DataFile Version 4.2\n");
57    header.push_str("trx-rs tractogram SPACE=RAS\n");
58    header.push_str("BINARY\n");
59    header.push_str("DATASET POLYDATA\n");
60    header.push_str(&format!("POINTS {} float\n", tractogram.nb_vertices()));
61    file.write_all(header.as_bytes())?;
62
63    for point in tractogram.positions() {
64        file.write_all(&point[0].to_be_bytes())?;
65        file.write_all(&point[1].to_be_bytes())?;
66        file.write_all(&point[2].to_be_bytes())?;
67    }
68
69    let total_line_entries: usize = tractogram
70        .streamlines()
71        .map(|streamline| streamline.len() + 1)
72        .sum();
73
74    let lines_header = format!(
75        "\nLINES {} {}\n",
76        tractogram.nb_streamlines(),
77        total_line_entries
78    );
79    file.write_all(lines_header.as_bytes())?;
80
81    let mut vertex_index = 0u32;
82    for streamline in tractogram.streamlines() {
83        let len = streamline.len() as u32;
84        file.write_all(&len.to_be_bytes())?;
85        for _ in 0..len {
86            file.write_all(&vertex_index.to_be_bytes())?;
87            vertex_index += 1;
88        }
89    }
90
91    file.flush()?;
92
93    Ok(())
94}
95
96pub fn inspect_vtk_declared_space(path: &Path) -> Result<Option<VtkCoordinateSpace>> {
97    let bytes = std::fs::read(path)?;
98    inspect_vtk_declared_space_bytes(&bytes)
99}
100
101pub fn vtk_import_warnings(path: &Path, mode: VtkCoordinateMode) -> Result<Vec<String>> {
102    let declared = inspect_vtk_declared_space(path)?;
103    Ok(match (mode, declared) {
104        (VtkCoordinateMode::HeaderOrWarn, None) => vec![
105            "VTK file does not declare `SPACE=RAS` or `SPACE=LPS`; trx-rs assumed LPS and converted to RAS.".to_string(),
106            "If the tractogram looks mirrored, re-import it and force VTK coordinates to RAS.".to_string(),
107        ],
108        _ => Vec::new(),
109    })
110}
111
112fn parse_vtk_bytes(
113    bytes: &[u8],
114    coordinate_mode: VtkCoordinateMode,
115) -> Result<(Vec<[f32; 3]>, Vec<u32>)> {
116    let mut cursor = 0usize;
117    let version = read_line(bytes, &mut cursor)?;
118    if !version.starts_with("# vtk DataFile Version") {
119        return Err(TrxError::Format("not a legacy VTK file".into()));
120    }
121
122    let _comment = read_line(bytes, &mut cursor)?;
123    let format = read_line(bytes, &mut cursor)?;
124    let dataset = read_line(bytes, &mut cursor)?;
125    if dataset.trim() != "DATASET POLYDATA" {
126        return Err(TrxError::Format(
127            "only VTK POLYDATA streamline files are supported".into(),
128        ));
129    }
130
131    let points_header = read_line(bytes, &mut cursor)?;
132    let mut header_parts = points_header.split_whitespace();
133    if header_parts.next() != Some("POINTS") {
134        return Err(TrxError::Format(
135            "VTK file is missing POINTS section".into(),
136        ));
137    }
138    let point_count = header_parts
139        .next()
140        .ok_or_else(|| TrxError::Format("missing VTK point count".into()))?
141        .parse::<usize>()
142        .map_err(|_| TrxError::Format("invalid VTK point count".into()))?;
143    let point_type = header_parts
144        .next()
145        .ok_or_else(|| TrxError::Format("missing VTK points datatype".into()))?;
146    let vtk_header_text = header_text(bytes, cursor)?;
147
148    let is_binary = format.trim().eq_ignore_ascii_case("BINARY");
149    let (mut positions, offsets) = if is_binary {
150        parse_binary_points(bytes, &mut cursor, point_count, point_type)?
151    } else if format.trim().eq_ignore_ascii_case("ASCII") {
152        parse_ascii_points_and_lines(bytes, cursor, point_count)?
153    } else {
154        return Err(TrxError::Format(format!(
155            "unsupported VTK encoding '{}'",
156            format.trim()
157        )));
158    };
159
160    let coordinate_space = resolve_vtk_coordinate_space(vtk_header_text, coordinate_mode)?;
161    for point in &mut positions {
162        *point = vtk_world_to_ras(*point, coordinate_space);
163    }
164
165    Ok((positions, offsets))
166}
167
168fn parse_ascii_points_and_lines(
169    bytes: &[u8],
170    cursor: usize,
171    point_count: usize,
172) -> Result<(Vec<[f32; 3]>, Vec<u32>)> {
173    let text = std::str::from_utf8(&bytes[cursor..])
174        .map_err(|_| TrxError::Format("ASCII VTK body is not valid UTF-8".into()))?;
175    let mut tokens = text.split_whitespace();
176
177    let mut points = Vec::with_capacity(point_count);
178    for _ in 0..point_count {
179        let x = next_token(&mut tokens, "point x")?
180            .parse::<f32>()
181            .map_err(|_| TrxError::Format("invalid VTK point coordinate".into()))?;
182        let y = next_token(&mut tokens, "point y")?
183            .parse::<f32>()
184            .map_err(|_| TrxError::Format("invalid VTK point coordinate".into()))?;
185        let z = next_token(&mut tokens, "point z")?
186            .parse::<f32>()
187            .map_err(|_| TrxError::Format("invalid VTK point coordinate".into()))?;
188        points.push([x, y, z]);
189    }
190
191    if next_token(&mut tokens, "LINES keyword")? != "LINES" {
192        return Err(TrxError::Format("VTK file is missing LINES section".into()));
193    }
194    let line_count = next_token(&mut tokens, "line count")?
195        .parse::<usize>()
196        .map_err(|_| TrxError::Format("invalid VTK line count".into()))?;
197    let _line_size = next_token(&mut tokens, "line size")?
198        .parse::<usize>()
199        .map_err(|_| TrxError::Format("invalid VTK line size".into()))?;
200
201    build_streamlines_from_tokens(points, line_count, &mut tokens)
202}
203
204fn parse_binary_points(
205    bytes: &[u8],
206    cursor: &mut usize,
207    point_count: usize,
208    point_type: &str,
209) -> Result<(Vec<[f32; 3]>, Vec<u32>)> {
210    let element_size = match point_type {
211        "float" => 4,
212        "double" => 8,
213        other => {
214            return Err(TrxError::Format(format!(
215                "unsupported binary VTK point datatype '{other}'"
216            )))
217        }
218    };
219    let points_bytes = point_count
220        .checked_mul(3)
221        .and_then(|count| count.checked_mul(element_size))
222        .ok_or_else(|| TrxError::Format("VTK points section is too large".into()))?;
223    let end = cursor
224        .checked_add(points_bytes)
225        .ok_or_else(|| TrxError::Format("VTK points section overflow".into()))?;
226    let data = bytes
227        .get(*cursor..end)
228        .ok_or_else(|| TrxError::Format("VTK points section is truncated".into()))?;
229
230    let mut points = Vec::with_capacity(point_count);
231    match point_type {
232        "float" => {
233            for chunk in data.as_chunks::<12>().0 {
234                points.push([
235                    f32::from_be_bytes(chunk[0..4].try_into().unwrap()),
236                    f32::from_be_bytes(chunk[4..8].try_into().unwrap()),
237                    f32::from_be_bytes(chunk[8..12].try_into().unwrap()),
238                ]);
239            }
240        }
241        "double" => {
242            for chunk in data.as_chunks::<24>().0 {
243                points.push([
244                    f64::from_be_bytes(chunk[0..8].try_into().unwrap()) as f32,
245                    f64::from_be_bytes(chunk[8..16].try_into().unwrap()) as f32,
246                    f64::from_be_bytes(chunk[16..24].try_into().unwrap()) as f32,
247                ]);
248            }
249        }
250        _ => unreachable!(),
251    }
252    *cursor = end;
253    skip_newlines(bytes, cursor);
254
255    let lines_header = read_line(bytes, cursor)?;
256    let mut parts = lines_header.split_whitespace();
257    if parts.next() != Some("LINES") {
258        return Err(TrxError::Format("VTK file is missing LINES section".into()));
259    }
260    let line_count = parts
261        .next()
262        .ok_or_else(|| TrxError::Format("missing VTK line count".into()))?
263        .parse::<usize>()
264        .map_err(|_| TrxError::Format("invalid VTK line count".into()))?;
265    let total_size = parts
266        .next()
267        .ok_or_else(|| TrxError::Format("missing VTK total line size".into()))?
268        .parse::<usize>()
269        .map_err(|_| TrxError::Format("invalid VTK total line size".into()))?;
270
271    let line_bytes = total_size
272        .checked_mul(4)
273        .ok_or_else(|| TrxError::Format("VTK line section is too large".into()))?;
274    let end = cursor
275        .checked_add(line_bytes)
276        .ok_or_else(|| TrxError::Format("VTK line section overflow".into()))?;
277    let data = bytes
278        .get(*cursor..end)
279        .ok_or_else(|| TrxError::Format("VTK line section is truncated".into()))?;
280
281    let mut ints = Vec::with_capacity(total_size);
282    for chunk in data.as_chunks::<4>().0 {
283        ints.push(i32::from_be_bytes(*chunk));
284    }
285
286    build_streamlines_from_ints(points, line_count, ints)
287}
288
289fn build_streamlines_from_ints(
290    points: Vec<[f32; 3]>,
291    line_count: usize,
292    ints: Vec<i32>,
293) -> Result<(Vec<[f32; 3]>, Vec<u32>)> {
294    let mut cursor = 0usize;
295    let mut positions = Vec::with_capacity(ints.len().saturating_sub(line_count));
296    let mut offsets = Vec::with_capacity(line_count + 1);
297    offsets.push(0);
298
299    for _ in 0..line_count {
300        let length = *ints
301            .get(cursor)
302            .ok_or_else(|| TrxError::Format("VTK LINES section ended unexpectedly".into()))?;
303        cursor += 1;
304        let length = usize::try_from(length)
305            .map_err(|_| TrxError::Format("VTK streamline length cannot be negative".into()))?;
306        for _ in 0..length {
307            let point_index = *ints.get(cursor).ok_or_else(|| {
308                TrxError::Format("VTK LINES point index section ended unexpectedly".into())
309            })?;
310            cursor += 1;
311            let point_index = usize::try_from(point_index)
312                .map_err(|_| TrxError::Format("VTK point index cannot be negative".into()))?;
313            let point = *points.get(point_index).ok_or_else(|| {
314                TrxError::Format(format!("VTK point index {point_index} is out of bounds"))
315            })?;
316            positions.push(point);
317        }
318        offsets.push(u32::try_from(positions.len()).map_err(|_| {
319            TrxError::Argument("tractogram has more than u32::MAX vertices".into())
320        })?);
321    }
322    Ok((positions, offsets))
323}
324
325fn build_streamlines_from_tokens<'a>(
326    points: Vec<[f32; 3]>,
327    line_count: usize,
328    tokens: &mut impl Iterator<Item = &'a str>,
329) -> Result<(Vec<[f32; 3]>, Vec<u32>)> {
330    let mut positions = Vec::new();
331    let mut offsets = Vec::with_capacity(line_count + 1);
332    offsets.push(0);
333
334    for _ in 0..line_count {
335        let length = next_token(tokens, "streamline length")?
336            .parse::<usize>()
337            .map_err(|_| TrxError::Format("invalid VTK streamline length".into()))?;
338        for _ in 0..length {
339            let point_index = next_token(tokens, "streamline point index")?
340                .parse::<usize>()
341                .map_err(|_| TrxError::Format("invalid VTK point index".into()))?;
342            let point = *points.get(point_index).ok_or_else(|| {
343                TrxError::Format(format!("VTK point index {point_index} is out of bounds"))
344            })?;
345            positions.push(point);
346        }
347        offsets.push(u32::try_from(positions.len()).map_err(|_| {
348            TrxError::Argument("tractogram has more than u32::MAX vertices".into())
349        })?);
350    }
351    Ok((positions, offsets))
352}
353
354fn read_line<'a>(bytes: &'a [u8], cursor: &mut usize) -> Result<&'a str> {
355    let start = *cursor;
356    while *cursor < bytes.len() && bytes[*cursor] != b'\n' {
357        *cursor += 1;
358    }
359    let end = *cursor;
360    if *cursor < bytes.len() && bytes[*cursor] == b'\n' {
361        *cursor += 1;
362    }
363    std::str::from_utf8(&bytes[start..end])
364        .map(|line| line.trim_end_matches('\r'))
365        .map_err(|_| TrxError::Format("VTK header is not valid UTF-8".into()))
366}
367
368fn skip_newlines(bytes: &[u8], cursor: &mut usize) {
369    while *cursor < bytes.len() && matches!(bytes[*cursor], b'\n' | b'\r') {
370        *cursor += 1;
371    }
372}
373
374fn next_token<'a>(tokens: &mut impl Iterator<Item = &'a str>, label: &str) -> Result<&'a str> {
375    tokens
376        .next()
377        .ok_or_else(|| TrxError::Format(format!("missing VTK token for {label}")))
378}
379
380fn inspect_vtk_declared_space_bytes(bytes: &[u8]) -> Result<Option<VtkCoordinateSpace>> {
381    let mut cursor = 0usize;
382    let version = read_line(bytes, &mut cursor)?;
383    if !version.starts_with("# vtk DataFile Version") {
384        return Err(TrxError::Format("not a legacy VTK file".into()));
385    }
386    let _comment = read_line(bytes, &mut cursor)?;
387    let _format = read_line(bytes, &mut cursor)?;
388    let _dataset = read_line(bytes, &mut cursor)?;
389    let _points = read_line(bytes, &mut cursor)?;
390    Ok(parse_vtk_declared_space(header_text(bytes, cursor)?))
391}
392
393fn parse_vtk_declared_space(comment: &str) -> Option<VtkCoordinateSpace> {
394    let upper = comment.to_ascii_uppercase();
395    if upper.contains("SPACE=RAS") {
396        Some(VtkCoordinateSpace::Ras)
397    } else if upper.contains("SPACE=LPS") {
398        Some(VtkCoordinateSpace::Lps)
399    } else {
400        None
401    }
402}
403
404fn resolve_vtk_coordinate_space(
405    header_text: &str,
406    mode: VtkCoordinateMode,
407) -> Result<VtkCoordinateSpace> {
408    Ok(match mode {
409        VtkCoordinateMode::AssumeRas => VtkCoordinateSpace::Ras,
410        VtkCoordinateMode::AssumeLps => VtkCoordinateSpace::Lps,
411        VtkCoordinateMode::HeaderOrWarn => {
412            parse_vtk_declared_space(header_text).unwrap_or(VtkCoordinateSpace::Lps)
413        }
414    })
415}
416
417fn header_text(bytes: &[u8], end: usize) -> Result<&str> {
418    std::str::from_utf8(&bytes[..end])
419        .map_err(|_| TrxError::Format("VTK header is not valid UTF-8".into()))
420}
421
422fn vtk_world_to_ras(point: [f32; 3], coordinate_space: VtkCoordinateSpace) -> [f32; 3] {
423    match coordinate_space {
424        VtkCoordinateSpace::Ras => point,
425        VtkCoordinateSpace::Lps => [-point[0], -point[1], point[2]],
426    }
427}
428
429#[cfg(test)]
430mod tests {
431    use super::{
432        parse_vtk_declared_space, resolve_vtk_coordinate_space, vtk_world_to_ras,
433        VtkCoordinateMode, VtkCoordinateSpace,
434    };
435
436    #[test]
437    fn vtk_header_space_parser_detects_ras_and_lps() {
438        assert_eq!(
439            parse_vtk_declared_space("created in slicer SPACE=RAS"),
440            Some(VtkCoordinateSpace::Ras)
441        );
442        assert_eq!(
443            parse_vtk_declared_space("created in slicer SPACE=LPS"),
444            Some(VtkCoordinateSpace::Lps)
445        );
446        assert_eq!(parse_vtk_declared_space("vtk output"), None);
447    }
448
449    #[test]
450    fn header_or_warn_defaults_to_lps_when_header_is_absent() {
451        assert_eq!(
452            resolve_vtk_coordinate_space("vtk output", VtkCoordinateMode::HeaderOrWarn).unwrap(),
453            VtkCoordinateSpace::Lps
454        );
455    }
456
457    #[test]
458    fn forced_ras_leaves_coordinates_unchanged() {
459        assert_eq!(
460            vtk_world_to_ras([1.0, 2.0, 3.0], VtkCoordinateSpace::Ras),
461            [1.0, 2.0, 3.0]
462        );
463    }
464}
trx_rs/formats/vtk.rs

trx_rs/formats/
vtk.rs
