Merge pull request #101 from qbio/feature/nifti-extension-writer

Feature/nifti extension writer (task 5 of issue #19)

Cargo.toml

@@ -47,6 +47,10 @@ tempfile = "3.1"
 name = "niftidump"
 path = "examples/niftidump/main.rs"
 
+[[example]]
+name = "gen_nifti"
+path = "examples/gen_nifti/main.rs"
+
 [features]
 default = ["ndarray_volumes"]
 nalgebra_affine = ["nalgebra", "simba"]

examples/gen_nifti/main.rs

@@ -0,0 +1,36 @@
+//! An application for writing a NIFTI file from scratch
+
+extern crate nifti;
+
+use std::env;
+
+#[cfg(feature = "ndarray_volumes")]
+use crate::nifti::{writer::WriterOptions, Extender, Extension, ExtensionSequence};
+
+#[cfg(feature = "ndarray_volumes")]
+fn main() {
+    let mut args = env::args().skip(1);
+    let filename = args.next().expect("Path to NIFTI file is required");
+
+    // generate some test data 256x256 float32
+    let data = ndarray::Array3::<f32>::zeros((256, 256, 1));
+
+    let extension1 = Extension::new(8 + 4, 3, vec![0, 0, 0, 0]);
+
+    let extension2 = Extension::from_str(6, "Hello World!");
+
+    let extension_sequence = ExtensionSequence::new(
+        Extender::from([1u8, 0u8, 0u8, 0u8]),
+        vec![extension1, extension2],
+    );
+
+    WriterOptions::new(&filename)
+        .with_extensions(extension_sequence)
+        .write_nifti(&data)
+        .unwrap();
+}
+
+#[cfg(not(feature = "ndarray_volumes"))]
+fn main() {
+    println!("This example requires the ndarray_volumes feature to be enabled");
+}

src/extension.rs

@@ -6,8 +6,31 @@
 
 use crate::error::{NiftiError, Result};
 use byteordered::{ByteOrdered, Endian};
+use num_derive::FromPrimitive;
 use std::io::{ErrorKind as IoErrorKind, Read};
 
+/// Data type for representing a NIfTI-1.1 extension code
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, FromPrimitive)]
+#[repr(u32)]
+pub enum NiftiEcode {
+    /// Ignore the extension
+    NiftEcodeIgnore = 0,
+    /// DICOM
+    NiftiEcodeDicom = 2,
+    /// AFNI extension in XML format
+    NiftiEcodeAFNI = 4,
+    /// String Comment
+    NiftiEcodeComment = 6,
+    /// XCEDE extension in XML format
+    NiftiEcodeXCEDE = 8,
+    /// JimDimInfo
+    NiftiEcodeJimDimInfo = 10,
+    /// WorkflowFWDS
+    NiftiEcodeWorkflowFWDS = 12,
+    /// Freesurfer
+    NiftiEcodeFreesurfer = 14,
+}
+
 /// Data type for the extender code.
 #[derive(Debug, Default, PartialEq, Clone, Copy)]
 pub struct Extender([u8; 4]);
@@ -82,6 +105,16 @@ impl Extension {
         }
     }
 
+    /// Create a new extension out of a &str
+    pub fn from_str(ecode: i32, edata: &str) -> Self {
+        let esize = 8 + edata.len() as i32;
+        // pad the esize to a multiple of 16
+        let padded_esize = (esize + 15) & !15;
+        let mut edata = edata.as_bytes().to_vec();
+        edata.resize(padded_esize as usize - 8, 0);
+        Extension::new(padded_esize, ecode, edata)
+    }
+
     /// Obtain the claimed extension raw size (`esize` field).
     pub fn size(&self) -> i32 {
         self.esize
@@ -130,6 +163,14 @@ impl<'a> IntoIterator for &'a ExtensionSequence {
 }
 
 impl ExtensionSequence {
+    /// Provide a public constructor
+    pub fn new(extender: Extender, extensions: Vec<Extension>) -> Self {
+        ExtensionSequence {
+            extender,
+            extensions,
+        }
+    }
+
     /// Read a sequence of extensions from a source, up until `len` bytes.
     pub fn from_reader<S, E>(
         extender: Extender,
@@ -186,6 +227,13 @@ impl ExtensionSequence {
         self.extensions.len()
     }
 
+    /// Return the number of bytes the extensions take on disk
+    pub fn bytes_on_disk(&self) -> usize {
+        self.extensions
+            .iter()
+            .map(|e| e.size() as usize)
+            .sum::<usize>()
+    }
     /// Get the extender code from this extension sequence.
     pub fn extender(&self) -> Extender {
         self.extender

src/header.rs

@@ -90,7 +90,10 @@ pub struct NiftiHeader {
     pub slice_start: i16,
     /// Grid spacings
     pub pixdim: [f32; 8],
-    /// Offset into .nii file to reach the volume
+    /// Offset into .nii file to reach the volume.
+    ///
+    /// Note: the highly unusual choice of f32 is intentional
+    /// and due to trying to achieve header backwards compatibility to ANALYZE 7.5
     pub vox_offset: f32,
     /// Data scaling: slope
     pub scl_slope: f32,

src/writer.rs

@@ -14,7 +14,7 @@ use crate::{
     header::{MAGIC_CODE_NI1, MAGIC_CODE_NIP1},
     util::{adapt_bytes, is_gz_file, is_hdr_file},
     volume::shape::Dim,
-    DataElement, NiftiHeader, NiftiType, Result,
+    DataElement, ExtensionSequence, NiftiHeader, NiftiType, Result,
 };
 
 #[derive(Debug, Clone, PartialEq)]
@@ -57,6 +57,9 @@ pub struct WriterOptions<'a> {
     /// The header file will only be compressed if the caller specifically asked for a path ending
     /// with "hdr.gz". Otherwise, only the volume will be compressed (if requested).
     force_header_compression: bool,
+
+    /// Optional ExtensionSequence
+    extension_sequence: Option<ExtensionSequence>,
 }
 
 impl<'a> WriterOptions<'a> {
@@ -81,6 +84,7 @@ impl<'a> WriterOptions<'a> {
             write_header_file,
             compression,
             force_header_compression: write_header_file && compression.is_some(),
+            extension_sequence: None,
         }
     }
 
@@ -127,6 +131,12 @@ impl<'a> WriterOptions<'a> {
         self
     }
 
+    /// Sets an extension sequence for the writer
+    pub fn with_extensions(mut self, extension_sequence: ExtensionSequence) -> Self {
+        self.extension_sequence = Some(extension_sequence);
+        self
+    }
+
     /// Write a nifti file (.nii or .nii.gz).
     pub fn write_nifti<A, S, D>(&self, data: &ArrayBase<S, D>) -> Result<()>
     where
@@ -149,12 +159,14 @@ impl<'a> WriterOptions<'a> {
                     header.endianness,
                 );
                 write_header(writer.as_mut(), &header)?;
+                write_extensions(writer.as_mut(), self.extension_sequence.as_ref())?;
                 write_data::<_, A, _, _, _, _>(writer.as_mut(), data)?;
                 let _ = writer.into_inner().finish()?;
             } else {
                 let mut writer =
                     ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
                 write_header(writer.as_mut(), &header)?;
+                write_extensions(writer.as_mut(), self.extension_sequence.as_ref())?;
                 write_data::<_, A, _, _, _, _>(writer, data)?;
             }
         } else {
@@ -165,6 +177,7 @@ impl<'a> WriterOptions<'a> {
                     header.endianness,
                 );
                 write_header(writer.as_mut(), &header)?;
+                write_extensions(writer.as_mut(), self.extension_sequence.as_ref())?;
                 let _ = writer.into_inner().finish()?;
 
                 let mut writer = ByteOrdered::runtime(
@@ -174,9 +187,10 @@ impl<'a> WriterOptions<'a> {
                 write_data::<_, A, _, _, _, _>(writer.as_mut(), data)?;
                 let _ = writer.into_inner().finish()?;
             } else {
-                let header_writer =
+                let mut header_writer =
                     ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
-                write_header(header_writer, &header)?;
+                write_header(header_writer.as_mut(), &header)?;
+                write_extensions(header_writer.as_mut(), self.extension_sequence.as_ref())?;
                 let data_writer =
                     ByteOrdered::runtime(BufWriter::new(data_file), header.endianness);
                 write_data::<_, A, _, _, _, _>(data_writer, data)?;
@@ -206,12 +220,14 @@ impl<'a> WriterOptions<'a> {
                     header.endianness,
                 );
                 write_header(writer.as_mut(), &header)?;
+                write_extensions(writer.as_mut(), self.extension_sequence.as_ref())?;
                 write_data::<_, u8, _, _, _, _>(writer.as_mut(), data)?;
                 let _ = writer.into_inner().finish()?;
             } else {
                 let mut writer =
                     ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
                 write_header(writer.as_mut(), &header)?;
+                write_extensions(writer.as_mut(), self.extension_sequence.as_ref())?;
                 write_data::<_, u8, _, _, _, _>(writer, data)?;
             }
         } else {
@@ -222,6 +238,7 @@ impl<'a> WriterOptions<'a> {
                     header.endianness,
                 );
                 write_header(writer.as_mut(), &header)?;
+                write_extensions(writer.as_mut(), self.extension_sequence.as_ref())?;
                 let _ = writer.into_inner().finish()?;
 
                 let mut writer = ByteOrdered::runtime(
@@ -231,10 +248,10 @@ impl<'a> WriterOptions<'a> {
                 write_data::<_, u8, _, _, _, _>(writer.as_mut(), data)?;
                 let _ = writer.into_inner().finish()?;
             } else {
-                let header_writer =
+                let mut header_writer =
                     ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
-                write_header(header_writer, &header)?;
-
+                write_header(header_writer.as_mut(), &header)?;
+                write_extensions(header_writer.as_mut(), self.extension_sequence.as_ref())?;
                 let data_writer =
                     ByteOrdered::runtime(BufWriter::new(data_file), header.endianness);
                 write_data::<_, u8, _, _, _, _>(data_writer, data)?;
@@ -253,12 +270,18 @@ impl<'a> WriterOptions<'a> {
         T: Data,
         D: Dimension,
     {
+        let mut vox_offset: f32 = 352.0;
+
+        if let Some(extension_sequence) = self.extension_sequence.as_ref() {
+            vox_offset += extension_sequence.bytes_on_disk() as f32;
+        }
+
         let mut header = NiftiHeader {
             dim: *Dim::from_slice(data.shape())?.raw(),
             sizeof_hdr: 348,
             datatype: datatype as i16,
             bitpix: (datatype.size_of() * 8) as i16,
-            vox_offset: 352.0,
+            vox_offset,
             scl_inter: 0.0,
             scl_slope: 1.0,
             magic: *MAGIC_CODE_NIP1,
@@ -297,6 +320,37 @@ impl<'a> WriterOptions<'a> {
     }
 }
 
+fn write_extensions<W, E>(
+    mut writer: ByteOrdered<W, E>,
+    extensions: Option<&ExtensionSequence>,
+) -> Result<()>
+where
+    W: Write,
+    E: Endian,
+{
+    let extensions = match extensions {
+        Some(extensions) => extensions,
+        None => {
+            writer.write_u32(0)?;
+            return Ok(());
+        }
+    };
+
+    if extensions.is_empty() {
+        // Write an extender code of 4 zeros, which for NIFTI means that there are no extensions
+        writer.write_u32(0)?;
+        return Ok(());
+    }
+
+    writer.write_all(extensions.extender().as_bytes())?;
+    for extension in extensions.iter() {
+        writer.write_i32(extension.size())?;
+        writer.write_i32(extension.code())?;
+        writer.write_all(extension.data())?;
+    }
+    Ok(())
+}
+
 fn write_header<W, E>(mut writer: ByteOrdered<W, E>, header: &NiftiHeader) -> Result<()>
 where
     W: Write,
@@ -360,10 +414,6 @@ where
     writer.write_all(&header.intent_name)?;
     writer.write_all(&header.magic)?;
 
-    // Empty 4 bytes after the header
-    // TODO(#19) Support writing extension data.
-    writer.write_u32(0)?;
-
     Ok(())
 }
 

tests/writer.rs

@@ -407,4 +407,28 @@ mod tests {
             fs::read("resources/rgb/4D.nii").unwrap()
         );
     }
+
+    #[test]
+    fn write_extended_header() {
+        let data: Array2<f64> = Array2::zeros((8, 8));
+
+        let path = get_temporary_path("2d_extended_header.nii");
+        let extension = nifti::Extension::from_str(6, "Hello World!");
+
+        let extension_sequence = nifti::ExtensionSequence::new(
+            nifti::Extender::from([1u8, 0u8, 0u8, 0u8]),
+            vec![extension],
+        );
+
+        WriterOptions::new(&path)
+            .with_extensions(extension_sequence)
+            .write_nifti(&data)
+            .unwrap();
+
+        // Verify the binary identity to the nibabel generated file
+        assert_eq!(
+            fs::read(&path).unwrap(),
+            fs::read("resources/minimal_extended_hdr.nii").unwrap()
+        );
+    }
 }