1use std::fs::File;
2use std::io::{Read, Write};
3use std::path::Path;
4
5use crate::Tractogram;
6
7pub fn load_trk(path: &Path) -> Result<Tractogram, Box<dyn std::error::Error>> {
8 let mut f = File::open(path)?;
9 let mut buffer = Vec::new();
10 f.read_to_end(&mut buffer)?;
11
12 if buffer.len() < 1000 {
13 return Err("File too small".into());
14 }
15
16 let n_scalars = i16::from_le_bytes(buffer[36..38].try_into().unwrap());
17 let n_properties = i16::from_le_bytes(buffer[238..240].try_into().unwrap());
18
19 let mut voxel_sizes = [
20 f32::from_le_bytes(buffer[12..16].try_into().unwrap()),
21 f32::from_le_bytes(buffer[16..20].try_into().unwrap()),
22 f32::from_le_bytes(buffer[20..24].try_into().unwrap()),
23 ];
24
25 if voxel_sizes[0] == 0.0 {
27 voxel_sizes[0] = 1.0;
28 }
29 if voxel_sizes[1] == 0.0 {
30 voxel_sizes[1] = 1.0;
31 }
32 if voxel_sizes[2] == 0.0 {
33 voxel_sizes[2] = 1.0;
34 }
35
36 let mut vox_to_ras = nalgebra::Matrix4::zeros();
37 let mut mat_offset = 440;
38 for r in 0..4 {
39 for c in 0..4 {
40 vox_to_ras[(r, c)] =
41 f32::from_le_bytes(buffer[mat_offset..mat_offset + 4].try_into().unwrap());
42 mat_offset += 4;
43 }
44 }
45
46 let mut tr = Tractogram::new();
47 let mut offset = 1000;
48
49 while offset + 4 <= buffer.len() {
50 let n_points = i32::from_le_bytes(buffer[offset..offset + 4].try_into().unwrap());
51 offset += 4;
52
53 if n_points < 0 {
54 return Err("Negative number of points in streamline".into());
55 }
56
57 let required_bytes = (n_points as usize) * (3 + n_scalars as usize) * 4;
58 if offset + required_bytes > buffer.len() {
59 return Err("Unexpected EOF reading streamline points".into());
60 }
61
62 let mut streamline = Vec::with_capacity(n_points as usize);
63 for _ in 0..n_points {
64 let raw_x = f32::from_le_bytes(buffer[offset..offset + 4].try_into().unwrap());
65 let raw_y = f32::from_le_bytes(buffer[offset + 4..offset + 8].try_into().unwrap());
66 let raw_z = f32::from_le_bytes(buffer[offset + 8..offset + 12].try_into().unwrap());
67
68 let cx = (raw_x / voxel_sizes[0]) - 0.5;
69 let cy = (raw_y / voxel_sizes[1]) - 0.5;
70 let cz = (raw_z / voxel_sizes[2]) - 0.5;
71
72 let p_vox = nalgebra::Point3::new(cx, cy, cz);
73 let p_ras = vox_to_ras.transform_point(&p_vox);
74
75 streamline.push([p_ras.x, p_ras.y, p_ras.z]);
76 offset += (3 + n_scalars as usize) * 4;
77 }
78 tr.push_streamline(&streamline)?;
79 offset += (n_properties as usize) * 4;
80 }
81
82 Ok(tr)
83}
84
85pub fn load_vtk(path: &Path) -> Result<Tractogram, Box<dyn std::error::Error>> {
86 let mut f = File::open(path)?;
87 let mut buffer = Vec::new();
88 f.read_to_end(&mut buffer)?;
89
90 let header_str = String::from_utf8_lossy(&buffer[0..std::cmp::min(1024, buffer.len())]);
91 let points_idx = header_str.find("POINTS ").ok_or("No POINTS")?;
92
93 let points_str = header_str[points_idx..]
94 .split_whitespace()
95 .nth(1)
96 .ok_or("No POINTS count")?;
97 let num_points: usize = points_str.parse()?;
98
99 let mut is_double = false;
100 if let Some(type_str) = header_str[points_idx..].split_whitespace().nth(2) {
101 if type_str == "double" {
102 is_double = true;
103 }
104 }
105
106 let header_end = header_str[points_idx..]
107 .find('\n')
108 .ok_or("No newline after POINTS")?
109 + points_idx
110 + 1;
111 let mut pts = Vec::with_capacity(num_points * 3);
112
113 let mut offset = header_end;
114 for _ in 0..num_points * 3 {
115 if is_double {
116 let chunk = buffer
117 .get(offset..offset + 8)
118 .ok_or("Unexpected EOF reading points")?;
119 let val = f64::from_be_bytes(chunk.try_into().unwrap());
120 pts.push(val as f32);
121 offset += 8;
122 } else {
123 let chunk = buffer
124 .get(offset..offset + 4)
125 .ok_or("Unexpected EOF reading points")?;
126 let val = f32::from_be_bytes(chunk.try_into().unwrap());
127 pts.push(val);
128 offset += 4;
129 }
130 }
131
132 let search_window = std::cmp::min(offset + 1024, buffer.len());
133 let lines_str_chunk = String::from_utf8_lossy(&buffer[offset..search_window]);
134
135 let lines_idx_in_chunk = lines_str_chunk.find("LINES ").ok_or("No LINES")?;
136 let lines_idx = offset + lines_idx_in_chunk;
137
138 let lines_str = lines_str_chunk[lines_idx_in_chunk..]
139 .split_whitespace()
140 .nth(1)
141 .ok_or("No LINES count")?;
142 let num_lines: usize = lines_str.parse()?;
143
144 let lines_header_end = lines_str_chunk[lines_idx_in_chunk..]
145 .find('\n')
146 .ok_or("No newline after LINES")?
147 + lines_idx
148 + 1;
149 offset = lines_header_end;
150
151 let mut tr = Tractogram::new();
152
153 if buffer
154 .get(offset..)
155 .is_some_and(|b| b.starts_with(b"OFFSETS"))
156 {
157 let offsets_header_end = buffer[offset..]
158 .iter()
159 .position(|&c| c == b'\n')
160 .ok_or("No newline after OFFSETS")?
161 + offset
162 + 1;
163 let is_int64 = buffer[offset..offsets_header_end]
164 .windows(5)
165 .any(|w| w == b"int64");
166 offset = offsets_header_end;
167
168 let mut offsets_vec = Vec::with_capacity(num_lines);
169 for _ in 0..num_lines {
170 if is_int64 {
171 let chunk = buffer
172 .get(offset..offset + 8)
173 .ok_or("Unexpected EOF reading offsets")?;
174 let val = u64::from_be_bytes(chunk.try_into().unwrap());
175 offsets_vec.push(val as usize);
176 offset += 8;
177 } else {
178 let chunk = buffer
179 .get(offset..offset + 4)
180 .ok_or("Unexpected EOF reading offsets")?;
181 let val = u32::from_be_bytes(chunk.try_into().unwrap());
182 offsets_vec.push(val as usize);
183 offset += 4;
184 }
185 }
186
187 for i in 0..num_lines - 1 {
188 let start = offsets_vec[i];
189 let end = offsets_vec[i + 1];
190 if end > pts.len() / 3 {
191 return Err("Offset points out of bounds".into());
192 }
193 let mut streamline = Vec::with_capacity(end.saturating_sub(start));
194 for pt_idx in start..end {
195 streamline.push([pts[pt_idx * 3], pts[pt_idx * 3 + 1], pts[pt_idx * 3 + 2]]);
196 }
197 tr.push_streamline(&streamline)?;
198 }
199 return Ok(tr);
200 }
201
202 let mut pt_idx = 0;
203 for _ in 0..num_lines {
204 if offset + 4 > buffer.len() {
205 break;
206 }
207 let n_pts = i32::from_be_bytes(buffer[offset..offset + 4].try_into().unwrap());
208 offset += 4;
209
210 if n_pts <= 0 {
211 continue;
212 }
213 if pt_idx + (n_pts as usize) > num_points {
214 break;
215 }
216
217 let mut streamline = Vec::with_capacity(n_pts as usize);
218 for _ in 0..n_pts {
219 offset += 4;
220 streamline.push([pts[pt_idx * 3], pts[pt_idx * 3 + 1], pts[pt_idx * 3 + 2]]);
221 pt_idx += 1;
222 }
223 tr.push_streamline(&streamline)?;
224 }
225
226 Ok(tr)
227}
228
229pub fn load_nifti_header(path: &Path) -> Result<crate::header::Header, Box<dyn std::error::Error>> {
230 let mut f = File::open(path)?;
231 let mut buffer = Vec::new();
232 f.read_to_end(&mut buffer)?;
233
234 if buffer.len() < 348 {
235 return Err("NIfTI file too small".into());
236 }
237
238 let mut sizeof_hdr_bytes = [0u8; 4];
239 sizeof_hdr_bytes.copy_from_slice(&buffer[0..4]);
240 let sizeof_hdr = i32::from_le_bytes(sizeof_hdr_bytes);
241
242 let is_nifti2;
243 let swap_endian;
244
245 if sizeof_hdr == 348 {
246 is_nifti2 = false;
247 swap_endian = false;
248 } else if sizeof_hdr == 348i32.swap_bytes() {
249 is_nifti2 = false;
250 swap_endian = true;
251 } else if sizeof_hdr == 540 {
252 is_nifti2 = true;
253 swap_endian = false;
254 } else if sizeof_hdr == 540i32.swap_bytes() {
255 is_nifti2 = true;
256 swap_endian = true;
257 } else {
258 return Err(format!("Unsupported NIfTI sizeof_hdr: {}", sizeof_hdr).into());
259 }
260
261 if is_nifti2 && buffer.len() < 540 {
262 return Err("NIfTI-2 file too small".into());
263 }
264
265 let read_i16 = |offset: usize| -> Result<i16, Box<dyn std::error::Error>> {
266 let bytes = buffer.get(offset..offset + 2).ok_or("Buffer too small")?;
267 let mut arr = [0u8; 2];
268 arr.copy_from_slice(bytes);
269 let val = if swap_endian {
270 i16::from_be_bytes(arr)
271 } else {
272 i16::from_le_bytes(arr)
273 };
274 Ok(val)
275 };
276
277 let read_i32 = |offset: usize| -> Result<i32, Box<dyn std::error::Error>> {
278 let bytes = buffer.get(offset..offset + 4).ok_or("Buffer too small")?;
279 let mut arr = [0u8; 4];
280 arr.copy_from_slice(bytes);
281 let val = if swap_endian {
282 i32::from_be_bytes(arr)
283 } else {
284 i32::from_le_bytes(arr)
285 };
286 Ok(val)
287 };
288
289 let read_i64 = |offset: usize| -> Result<i64, Box<dyn std::error::Error>> {
290 let bytes = buffer.get(offset..offset + 8).ok_or("Buffer too small")?;
291 let mut arr = [0u8; 8];
292 arr.copy_from_slice(bytes);
293 let val = if swap_endian {
294 i64::from_be_bytes(arr)
295 } else {
296 i64::from_le_bytes(arr)
297 };
298 Ok(val)
299 };
300
301 let read_f32 = |offset: usize| -> Result<f32, Box<dyn std::error::Error>> {
302 let bytes = buffer.get(offset..offset + 4).ok_or("Buffer too small")?;
303 let mut arr = [0u8; 4];
304 arr.copy_from_slice(bytes);
305 let val = if swap_endian {
306 f32::from_be_bytes(arr)
307 } else {
308 f32::from_le_bytes(arr)
309 };
310 Ok(val)
311 };
312
313 let read_f64 = |offset: usize| -> Result<f64, Box<dyn std::error::Error>> {
314 let bytes = buffer.get(offset..offset + 8).ok_or("Buffer too small")?;
315 let mut arr = [0u8; 8];
316 arr.copy_from_slice(bytes);
317 let val = if swap_endian {
318 f64::from_be_bytes(arr)
319 } else {
320 f64::from_le_bytes(arr)
321 };
322 Ok(val)
323 };
324
325 let mut dimensions = [1, 1, 1];
326 let qform_code;
327 let sform_code;
328
329 let mut pixdim = [1.0; 8];
330 let mut srow_x = [0.0; 4];
331 let mut srow_y = [0.0; 4];
332 let mut srow_z = [0.0; 4];
333 let quatern_b;
334 let quatern_c;
335 let quatern_d;
336 let qoffset_x;
337 let qoffset_y;
338 let qoffset_z;
339
340 if is_nifti2 {
341 for i in 1..=3 {
342 dimensions[i - 1] = read_i64(16 + i * 8)? as u64;
343 }
344 for (i, px) in pixdim.iter_mut().enumerate() {
345 *px = read_f64(80 + i * 8)?;
346 }
347 qform_code = read_i32(344)?;
348 sform_code = read_i32(348)?;
349 quatern_b = read_f64(352)?;
350 quatern_c = read_f64(360)?;
351 quatern_d = read_f64(368)?;
352 qoffset_x = read_f64(376)?;
353 qoffset_y = read_f64(384)?;
354 qoffset_z = read_f64(392)?;
355 for i in 0..4 {
356 srow_x[i] = read_f64(400 + i * 8)?;
357 srow_y[i] = read_f64(432 + i * 8)?;
358 srow_z[i] = read_f64(464 + i * 8)?;
359 }
360 } else {
361 for i in 1..=3 {
362 dimensions[i - 1] = read_i16(40 + i * 2)? as u64;
363 }
364 for (i, px) in pixdim.iter_mut().enumerate() {
365 *px = read_f32(76 + i * 4)? as f64;
366 }
367 qform_code = read_i16(252)? as i32;
368 sform_code = read_i16(254)? as i32;
369 quatern_b = read_f32(256)? as f64;
370 quatern_c = read_f32(260)? as f64;
371 quatern_d = read_f32(264)? as f64;
372 qoffset_x = read_f32(268)? as f64;
373 qoffset_y = read_f32(272)? as f64;
374 qoffset_z = read_f32(276)? as f64;
375 for i in 0..4 {
376 srow_x[i] = read_f32(280 + i * 4)? as f64;
377 srow_y[i] = read_f32(296 + i * 4)? as f64;
378 srow_z[i] = read_f32(312 + i * 4)? as f64;
379 }
380 }
381
382 let mut voxel_to_rasmm = crate::header::Header::identity_affine();
383
384 if sform_code > 0 {
385 voxel_to_rasmm[0] = srow_x;
386 voxel_to_rasmm[1] = srow_y;
387 voxel_to_rasmm[2] = srow_z;
388 voxel_to_rasmm[3] = [0.0, 0.0, 0.0, 1.0];
389 } else if qform_code > 0 {
390 let b = quatern_b;
391 let c = quatern_c;
392 let d = quatern_d;
393 let a = (1.0 - b * b - c * c - d * d).max(0.0).sqrt();
394 let qfac = if pixdim[0] == 0.0 { 1.0 } else { pixdim[0] };
395 let dx = pixdim[1];
396 let dy = pixdim[2];
397 let dz = pixdim[3];
398
399 let r00 = a * a + b * b - c * c - d * d;
400 let r01 = 2.0 * (b * c - a * d);
401 let r02 = 2.0 * (b * d + a * c);
402
403 let r10 = 2.0 * (b * c + a * d);
404 let r11 = a * a + c * c - b * b - d * d;
405 let r12 = 2.0 * (c * d - a * b);
406
407 let r20 = 2.0 * (b * d - a * c);
408 let r21 = 2.0 * (c * d + a * b);
409 let r22 = a * a + d * d - c * c - b * b;
410
411 voxel_to_rasmm[0] = [r00 * dx, r01 * dy, r02 * qfac * dz, qoffset_x];
412 voxel_to_rasmm[1] = [r10 * dx, r11 * dy, r12 * qfac * dz, qoffset_y];
413 voxel_to_rasmm[2] = [r20 * dx, r21 * dy, r22 * qfac * dz, qoffset_z];
414 voxel_to_rasmm[3] = [0.0, 0.0, 0.0, 1.0];
415 } else {
416 return Err("NIfTI file has no valid spatial transform".into());
417 }
418
419 let header = crate::header::Header {
420 voxel_to_rasmm,
421 dimensions,
422 nb_streamlines: 0,
423 nb_vertices: 0,
424 extra: Default::default(),
425 };
426 Ok(header)
427}
428
429pub fn write_trx(
430 path: &Path,
431 tractogram: &Tractogram,
432 ref_nifti: Option<&Path>,
433) -> Result<(), Box<dyn std::error::Error>> {
434 let mut tractogram = tractogram.clone();
435 let header_empty = tractogram.header().voxel_to_rasmm
436 == crate::header::Header::identity_affine()
437 && tractogram.header().dimensions == [1, 1, 1];
438
439 if header_empty {
440 if let Some(p) = ref_nifti {
441 let hdr = load_nifti_header(p)?;
442 tractogram.set_header(hdr);
443 } else {
444 return Err("TCK -> TRX requires a reference NIfTI file".into());
445 }
446 }
447
448 let any_trx = tractogram.to_trx(crate::dtype::DType::Float32)?;
449 any_trx.save(path)?;
450 Ok(())
451}
452
453pub fn write_trk(
454 path: &Path,
455 tractogram: &Tractogram,
456 ref_nifti: Option<&Path>,
457) -> Result<(), Box<dyn std::error::Error>> {
458 let mut file = File::create(path)?;
459 let mut header_bytes = vec![0u8; 1000];
460
461 header_bytes[0..5].copy_from_slice(b"TRACK");
462
463 let mut header = tractogram.header().clone();
464 let header_empty = header.voxel_to_rasmm == crate::header::Header::identity_affine()
465 && header.dimensions == [1, 1, 1];
466 if header_empty {
467 if let Some(p) = ref_nifti {
468 header = load_nifti_header(p)?;
469 } else {
470 return Err("TCK -> TRK requires a reference NIfTI file".into());
471 }
472 }
473 let dims = [
474 header.dimensions[0] as i16,
475 header.dimensions[1] as i16,
476 header.dimensions[2] as i16,
477 ];
478 header_bytes[6..8].copy_from_slice(&dims[0].to_le_bytes());
479 header_bytes[8..10].copy_from_slice(&dims[1].to_le_bytes());
480 header_bytes[10..12].copy_from_slice(&dims[2].to_le_bytes());
481
482 let vox_to_ras = header.voxel_to_rasmm;
483 let voxel_sizes = [
484 ((vox_to_ras[0][0].powi(2) + vox_to_ras[1][0].powi(2) + vox_to_ras[2][0].powi(2)).sqrt())
485 as f32,
486 ((vox_to_ras[0][1].powi(2) + vox_to_ras[1][1].powi(2) + vox_to_ras[2][1].powi(2)).sqrt())
487 as f32,
488 ((vox_to_ras[0][2].powi(2) + vox_to_ras[1][2].powi(2) + vox_to_ras[2][2].powi(2)).sqrt())
489 as f32,
490 ];
491 header_bytes[12..16].copy_from_slice(&voxel_sizes[0].to_le_bytes());
492 header_bytes[16..20].copy_from_slice(&voxel_sizes[1].to_le_bytes());
493 header_bytes[20..24].copy_from_slice(&voxel_sizes[2].to_le_bytes());
494
495 let mut offset = 440;
496 for row in &vox_to_ras {
497 for &elem in row {
498 let val = elem as f32;
499 header_bytes[offset..offset + 4].copy_from_slice(&val.to_le_bytes());
500 offset += 4;
501 }
502 }
503
504 header_bytes[948..952].copy_from_slice(b"RAS\0");
505
506 let nb_streamlines = tractogram.nb_streamlines() as i32;
507 header_bytes[988..992].copy_from_slice(&nb_streamlines.to_le_bytes());
508
509 header_bytes[992..996].copy_from_slice(&2i32.to_le_bytes());
510
511 header_bytes[996..1000].copy_from_slice(&1000i32.to_le_bytes());
512
513 file.write_all(&header_bytes)?;
514
515 let mut mat = nalgebra::Matrix4::zeros();
516 for r in 0..4 {
517 for c in 0..4 {
518 mat[(r, c)] = vox_to_ras[r][c] as f32;
519 }
520 }
521 let inv_mat = mat.try_inverse().unwrap_or(nalgebra::Matrix4::identity());
522
523 let offsets = tractogram.offsets();
524 let positions = tractogram.positions();
525 let mut chunk = Vec::with_capacity(4 * 1024 * 1024);
526
527 for i in 0..tractogram.nb_streamlines() {
528 let start = offsets[i] as usize;
529 let end = offsets[i + 1] as usize;
530 let n_points = (end - start) as i32;
531
532 chunk.extend_from_slice(&n_points.to_le_bytes());
533 for &pt in &positions[start..end] {
534 let p_ras = nalgebra::Point3::new(pt[0], pt[1], pt[2]);
535 let p_center = inv_mat.transform_point(&p_ras);
536
537 let vox_x = (p_center.x + 0.5) * voxel_sizes[0];
538 let vox_y = (p_center.y + 0.5) * voxel_sizes[1];
539 let vox_z = (p_center.z + 0.5) * voxel_sizes[2];
540
541 chunk.extend_from_slice(&vox_x.to_le_bytes());
542 chunk.extend_from_slice(&vox_y.to_le_bytes());
543 chunk.extend_from_slice(&vox_z.to_le_bytes());
544 }
545
546 if chunk.len() >= 4_000_000 {
547 file.write_all(&chunk)?;
548 chunk.clear();
549 }
550 }
551
552 if !chunk.is_empty() {
553 file.write_all(&chunk)?;
554 }
555
556 Ok(())
557}