Files
rustytorch/crates/specialized/rtx-medical-io/src/nifti/writer.rs
T
2026-03-04 00:08:42 +00:00

447 lines
12 KiB
Rust

//! NIfTI file writer supporting .nii and .nii.gz formats.
use crate::error::{MedicalIoError, Result};
use crate::nifti::header::{NiftiDataType, NiftiHeader, SpatialUnits, TransformCode};
use crate::volume::Volume;
use byteorder::{LittleEndian, WriteBytesExt};
use flate2::Compression;
use flate2::write::GzEncoder;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
/// Write a Volume to a NIfTI file.
///
/// The output format is determined by the file extension:
/// - `.nii` - uncompressed NIfTI-1
/// - `.nii.gz` - gzip-compressed NIfTI-1
///
/// # Arguments
/// * `volume` - The volume to write
/// * `path` - Output file path
/// * `datatype` - Output data type (default: Float64)
///
/// # Example
/// ```ignore
/// use rtx_medical_io::nifti::{write_nifti, NiftiDataType};
/// use rtx_medical_io::volume::Volume;
///
/// let volume = Volume::zeros([64, 64, 64]);
/// write_nifti(&volume, "output.nii.gz", NiftiDataType::Float32)?;
/// ```
pub fn write_nifti<P: AsRef<Path>>(
volume: &Volume,
path: P,
datatype: NiftiDataType,
) -> Result<()> {
let path = path.as_ref();
let path_str = path.to_string_lossy().to_lowercase();
let is_gzipped = path_str.ends_with(".nii.gz") || path_str.ends_with(".gz");
// Build header
let header = build_header(volume, datatype);
// Convert data to target type
let raw_data = convert_from_f64(volume.data(), datatype)?;
if is_gzipped {
write_nifti_gz(path, &header, &raw_data)
} else {
write_nifti_uncompressed(path, &header, &raw_data)
}
}
/// Write a gzip-compressed NIfTI file
fn write_nifti_gz<P: AsRef<Path>>(path: P, header: &NiftiHeader, data: &[u8]) -> Result<()> {
let file = File::create(path)?;
let writer = BufWriter::new(file);
let mut encoder = GzEncoder::new(writer, Compression::default());
write_header(&mut encoder, header)?;
encoder.write_all(data)?;
encoder.finish().map_err(|e| {
MedicalIoError::Compression(format!("Failed to finish gzip compression: {}", e))
})?;
Ok(())
}
/// Write an uncompressed NIfTI file
fn write_nifti_uncompressed<P: AsRef<Path>>(
path: P,
header: &NiftiHeader,
data: &[u8],
) -> Result<()> {
let file = File::create(path)?;
let mut writer = BufWriter::new(file);
write_header(&mut writer, header)?;
writer.write_all(data)?;
writer.flush()?;
Ok(())
}
/// Build a NIfTI-1 header from a Volume
fn build_header(volume: &Volume, datatype: NiftiDataType) -> NiftiHeader {
let [x, y, z] = volume.shape();
let [dx, dy, dz] = volume.spacing();
let affine = volume.affine();
NiftiHeader {
sizeof_hdr: 348,
datatype,
bitpix: (datatype.bytes_per_voxel() * 8) as i16,
dim: [3, x as i64, y as i64, z as i64, 1, 1, 1, 1],
intent_p1: 0.0,
intent_p2: 0.0,
intent_p3: 0.0,
intent_code: 0,
pixdim: [1.0, dx, dy, dz, 1.0, 1.0, 1.0, 1.0],
vox_offset: 352,
scl_slope: 1.0,
scl_inter: 0.0,
slice_code: 0,
xyzt_units: SpatialUnits::Millimeters as u8,
cal_max: 0.0,
cal_min: 0.0,
slice_duration: 0.0,
toffset: 0.0,
slice_start: 0,
slice_end: 0,
descrip: "Created by rtx-medical-io".to_string(),
aux_file: String::new(),
qform_code: TransformCode::Unknown,
sform_code: TransformCode::ScannerAnat,
quatern_b: 0.0,
quatern_c: 0.0,
quatern_d: 0.0,
qoffset_x: 0.0,
qoffset_y: 0.0,
qoffset_z: 0.0,
srow_x: affine[0],
srow_y: affine[1],
srow_z: affine[2],
intent_name: String::new(),
magic: *b"n+1\0\0\0\0\0",
is_nifti2: false,
little_endian: true,
}
}
/// Write NIfTI-1 header (348 bytes + 4 bytes padding = 352 bytes)
fn write_header<W: Write>(writer: &mut W, header: &NiftiHeader) -> Result<()> {
// sizeof_hdr (0-3)
writer.write_i32::<LittleEndian>(header.sizeof_hdr)?;
// data_type (unused, 4-13) - 10 bytes
writer.write_all(&[0u8; 10])?;
// db_name (unused, 14-31) - 18 bytes
writer.write_all(&[0u8; 18])?;
// extents (unused, 32-35)
writer.write_i32::<LittleEndian>(0)?;
// session_error (unused, 36-37)
writer.write_i16::<LittleEndian>(0)?;
// regular (unused, 38)
writer.write_all(&[0u8; 1])?;
// dim_info (unused, 39)
writer.write_all(&[0u8; 1])?;
// dim (40-55) - 8 x i16
for i in 0..8 {
writer.write_i16::<LittleEndian>(header.dim[i] as i16)?;
}
// intent_p1 (56-59)
writer.write_f32::<LittleEndian>(header.intent_p1 as f32)?;
// intent_p2 (60-63)
writer.write_f32::<LittleEndian>(header.intent_p2 as f32)?;
// intent_p3 (64-67)
writer.write_f32::<LittleEndian>(header.intent_p3 as f32)?;
// intent_code (68-69)
writer.write_i16::<LittleEndian>(header.intent_code)?;
// datatype (70-71)
writer.write_i16::<LittleEndian>(header.datatype as i16)?;
// bitpix (72-73)
writer.write_i16::<LittleEndian>(header.bitpix)?;
// slice_start (74-75)
writer.write_i16::<LittleEndian>(header.slice_start as i16)?;
// pixdim (76-107) - 8 x f32
for i in 0..8 {
writer.write_f32::<LittleEndian>(header.pixdim[i] as f32)?;
}
// vox_offset (108-111)
writer.write_f32::<LittleEndian>(header.vox_offset as f32)?;
// scl_slope (112-115)
writer.write_f32::<LittleEndian>(header.scl_slope as f32)?;
// scl_inter (116-119)
writer.write_f32::<LittleEndian>(header.scl_inter as f32)?;
// slice_end (120-121)
writer.write_i16::<LittleEndian>(header.slice_end as i16)?;
// slice_code (122)
writer.write_all(&[header.slice_code])?;
// xyzt_units (123)
writer.write_all(&[header.xyzt_units])?;
// cal_max (124-127)
writer.write_f32::<LittleEndian>(header.cal_max as f32)?;
// cal_min (128-131)
writer.write_f32::<LittleEndian>(header.cal_min as f32)?;
// slice_duration (132-135)
writer.write_f32::<LittleEndian>(header.slice_duration as f32)?;
// toffset (136-139)
writer.write_f32::<LittleEndian>(header.toffset as f32)?;
// glmax (unused, 140-143)
writer.write_i32::<LittleEndian>(0)?;
// glmin (unused, 144-147)
writer.write_i32::<LittleEndian>(0)?;
// descrip (148-227) - 80 bytes
let descrip_bytes = header.descrip.as_bytes();
let mut descrip_buf = [0u8; 80];
let len = descrip_bytes.len().min(80);
descrip_buf[..len].copy_from_slice(&descrip_bytes[..len]);
writer.write_all(&descrip_buf)?;
// aux_file (228-251) - 24 bytes
let aux_bytes = header.aux_file.as_bytes();
let mut aux_buf = [0u8; 24];
let len = aux_bytes.len().min(24);
aux_buf[..len].copy_from_slice(&aux_bytes[..len]);
writer.write_all(&aux_buf)?;
// qform_code (252-253)
writer.write_i16::<LittleEndian>(header.qform_code as i16)?;
// sform_code (254-255)
writer.write_i16::<LittleEndian>(header.sform_code as i16)?;
// quatern_b (256-259)
writer.write_f32::<LittleEndian>(header.quatern_b as f32)?;
// quatern_c (260-263)
writer.write_f32::<LittleEndian>(header.quatern_c as f32)?;
// quatern_d (264-267)
writer.write_f32::<LittleEndian>(header.quatern_d as f32)?;
// qoffset_x (268-271)
writer.write_f32::<LittleEndian>(header.qoffset_x as f32)?;
// qoffset_y (272-275)
writer.write_f32::<LittleEndian>(header.qoffset_y as f32)?;
// qoffset_z (276-279)
writer.write_f32::<LittleEndian>(header.qoffset_z as f32)?;
// srow_x (280-295) - 4 x f32
for val in &header.srow_x {
writer.write_f32::<LittleEndian>(*val as f32)?;
}
// srow_y (296-311) - 4 x f32
for val in &header.srow_y {
writer.write_f32::<LittleEndian>(*val as f32)?;
}
// srow_z (312-327) - 4 x f32
for val in &header.srow_z {
writer.write_f32::<LittleEndian>(*val as f32)?;
}
// intent_name (328-343) - 16 bytes
let intent_bytes = header.intent_name.as_bytes();
let mut intent_buf = [0u8; 16];
let len = intent_bytes.len().min(16);
intent_buf[..len].copy_from_slice(&intent_bytes[..len]);
writer.write_all(&intent_buf)?;
// magic (344-347) - 4 bytes for NIfTI-1
writer.write_all(&header.magic[..4])?;
// Padding to vox_offset (348-351) - 4 bytes
writer.write_all(&[0u8; 4])?;
Ok(())
}
/// Convert f64 data to raw bytes for the target data type
fn convert_from_f64(data: &[f64], datatype: NiftiDataType) -> Result<Vec<u8>> {
let mut result = Vec::with_capacity(data.len() * datatype.bytes_per_voxel());
match datatype {
NiftiDataType::UInt8 => {
for &val in data {
result.push(val.clamp(0.0, 255.0) as u8);
}
}
NiftiDataType::Int8 => {
for &val in data {
result.push(val.clamp(-128.0, 127.0) as i8 as u8);
}
}
NiftiDataType::Int16 => {
for &val in data {
let v = val.clamp(i16::MIN as f64, i16::MAX as f64) as i16;
result.write_i16::<LittleEndian>(v)?;
}
}
NiftiDataType::UInt16 => {
for &val in data {
let v = val.clamp(0.0, u16::MAX as f64) as u16;
result.write_u16::<LittleEndian>(v)?;
}
}
NiftiDataType::Int32 => {
for &val in data {
let v = val.clamp(i32::MIN as f64, i32::MAX as f64) as i32;
result.write_i32::<LittleEndian>(v)?;
}
}
NiftiDataType::UInt32 => {
for &val in data {
let v = val.clamp(0.0, u32::MAX as f64) as u32;
result.write_u32::<LittleEndian>(v)?;
}
}
NiftiDataType::Float32 => {
for &val in data {
result.write_f32::<LittleEndian>(val as f32)?;
}
}
NiftiDataType::Float64 => {
for &val in data {
result.write_f64::<LittleEndian>(val)?;
}
}
NiftiDataType::Int64 => {
for &val in data {
let v = val.clamp(i64::MIN as f64, i64::MAX as f64) as i64;
result.write_i64::<LittleEndian>(v)?;
}
}
NiftiDataType::UInt64 => {
for &val in data {
let v = val.clamp(0.0, u64::MAX as f64) as u64;
result.write_u64::<LittleEndian>(v)?;
}
}
_ => {
return Err(MedicalIoError::UnsupportedDataType(datatype as i16));
}
}
Ok(result)
}
/// Write a Volume to an INR file (INRIMAGE-4 format).
///
/// This format is used internally by CGAL for mesh generation.
/// It's a simple headerless format with a text header followed by raw binary data.
pub fn write_inr<P: AsRef<Path>>(volume: &Volume, path: P) -> Result<()> {
let path = path.as_ref();
let file = File::create(path)?;
let mut writer = BufWriter::new(file);
let [xdim, ydim, zdim] = volume.shape();
let [vx, vy, vz] = volume.spacing();
// Determine data type and bit length
let (btype, bitlen) = ("float", 64);
// Build header
let mut header = format!(
"#INRIMAGE-4#{{\n\
XDIM={}\n\
YDIM={}\n\
ZDIM={}\n\
VDIM=1\n\
TYPE={}\n\
PIXSIZE={} bits\n\
CPU=decm\n\
VX={:.6}\n\
VY={:.6}\n\
VZ={:.6}\n",
xdim, ydim, zdim, btype, bitlen, vx, vy, vz
);
// Pad header to 256 - 4 bytes, then add closing tag
let target_len = 256 - 4;
while header.len() < target_len {
header.push('\n');
}
header.push_str("##}\n");
// Write header
writer.write_all(header.as_bytes())?;
// Write data in Fortran order (column-major)
let data = volume.data();
for &val in data {
writer.write_f64::<LittleEndian>(val)?;
}
writer.flush()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_build_header() {
let volume = Volume::zeros([64, 64, 32]);
let header = build_header(&volume, NiftiDataType::Float32);
assert_eq!(header.sizeof_hdr, 348);
assert_eq!(header.dim[1], 64);
assert_eq!(header.dim[2], 64);
assert_eq!(header.dim[3], 32);
assert_eq!(header.datatype, NiftiDataType::Float32);
assert_eq!(header.bitpix, 32);
}
#[test]
fn test_convert_uint8() {
let data = vec![0.0, 127.5, 255.0, 300.0, -10.0];
let result = convert_from_f64(&data, NiftiDataType::UInt8).unwrap();
// 127.5 truncates to 127 (not rounded)
assert_eq!(result, vec![0, 127, 255, 255, 0]);
}
}