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

trx_rs/formats/
vtk.rs
