Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

RGB images writer #26

Merged
merged 10 commits into from
Nov 22, 2018
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Binary file added resources/rgb/3D.nii
Binary file not shown.
Binary file added resources/rgb/4D.nii
Binary file not shown.
302 changes: 128 additions & 174 deletions src/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,20 @@

use std::fs::File;
use std::io::BufWriter;
use std::mem;
use std::ops::{Div, Sub};
use std::path::Path;

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{Array, ArrayBase, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use ndarray::{ArrayBase, ArrayView, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {util::is_gz_file, volume::element::DataElement, NiftiHeader, Result};
use {
header::MAGIC_CODE_NIP1, util::is_gz_file, volume::element::DataElement, NiftiHeader,
NiftiType, Result,
};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
Expand Down Expand Up @@ -43,6 +45,71 @@ where
A: ScalarOperand,
A: Sub<Output = A>,
D: Dimension + RemoveAxis,
{
let header = build_header(data, reference, A::DATA_TYPE);

// Need the transpose for fortran ordering used in nifti file format.
let data = data.t();

let f = File::create(&path)?;
let mut writer = BufWriter::new(f);
if is_gz_file(&path) {
let mut e = GzEncoder::new(writer, Compression::fast());
write_header(&mut e, &header)?;
write_data(&mut e, &header, data)?;
let _ = e.finish()?; // Must use result
} else {
write_header(&mut writer, &header)?;
write_data(&mut writer, &header, data)?;
}
Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0. If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
path: P,
data: &ArrayBase<S, D>,
reference: Option<&NiftiHeader>,
) -> Result<()>
where
P: AsRef<Path>,
S: Data<Elem = [u8; 3]>,
D: Dimension + RemoveAxis,
{
// The `scl_slope` and `scl_inter` fields are ignored on the Rgb24 type.
let mut header = build_header(data, reference, NiftiType::Rgb24);
header.scl_slope = 1.0;
header.scl_inter = 0.0;

// Need the transpose for fortran used in nifti file format.
let data = data.t();

let f = File::create(&path)?;
let mut writer = BufWriter::new(f);
if is_gz_file(&path) {
let mut e = GzEncoder::new(writer, Compression::fast());
write_header(&mut e, &header)?;
write_slices(&mut e, data)?;
let _ = e.finish()?; // Must use result
} else {
write_header(&mut writer, &header)?;
write_slices(&mut writer, data)?;
}
Ok(())
}

fn build_header<T, D>(
data: &ArrayBase<T, D>,
reference: Option<&NiftiHeader>,
datatype: NiftiType,
) -> NiftiHeader
where
T: Data,
D: Dimension,
{
let mut dim = [1; 8];
dim[0] = data.ndim() as u16;
Expand All @@ -53,28 +120,27 @@ where
// If no reference header is given, use the default.
let reference = match reference {
Some(r) => r.clone(),
None => NiftiHeader::default(),
None => {
let mut header = NiftiHeader::default();
header.pixdim = [1.0; 8];
header.sform_code = 2;
header.srow_x = [1.0, 0.0, 0.0, 0.0];
header.srow_y = [0.0, 1.0, 0.0, 0.0];
header.srow_z = [0.0, 0.0, 1.0, 0.0];
header
}
};
let header = NiftiHeader {

NiftiHeader {
dim,
datatype: A::DATA_TYPE as i16,
bitpix: (mem::size_of::<A>() * 8) as i16,
sizeof_hdr: 348,
datatype: datatype as i16,
bitpix: (datatype.size_of() * 8) as i16,
vox_offset: 352.0,
magic: *MAGIC_CODE_NIP1,
// All other fields are copied from reference header
..reference
};

let f = File::create(&path)?;
let mut writer = BufWriter::new(f);
if is_gz_file(&path) {
let mut e = GzEncoder::new(writer, Compression::fast());
write_header(&mut e, &header)?;
write_data(&mut e, header, data)?;
let _ = e.finish()?; // Must use result
} else {
write_header(&mut writer, &header)?;
write_data(&mut writer, header, data)?;
}
Ok(())
}

fn write_header<W>(writer: &mut W, header: &NiftiHeader) -> Result<()>
Expand Down Expand Up @@ -149,179 +215,67 @@ where
/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<A, S, D, W>(writer: &mut W, header: NiftiHeader, data: &ArrayBase<S, D>) -> Result<()>
fn write_data<T, D, W>(writer: &mut W, header: &NiftiHeader, data: ArrayView<T, D>) -> Result<()>
where
S: Data<Elem = A>,
A: Copy,
A: DataElement,
A: Div<Output = A>,
A: FromPrimitive,
A: PodTransmutable,
A: ScalarOperand,
A: Sub<Output = A>,
T: Clone + PodTransmutable,
T: Div<Output = T>,
T: FromPrimitive,
T: PodTransmutable,
T: ScalarOperand,
T: Sub<Output = T>,
D: Dimension + RemoveAxis,
W: WriteBytesExt,
{
// Need the transpose for fortran ordering used in nifti file format.
let data = data.t();

let mut write_slice = |data: Array<A, D::Smaller>| -> Result<()> {
let len = data.len();
let arr_data = data.into_shape(len).unwrap();
let slice = arr_data.as_slice().unwrap();
writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
Ok(())
};

// `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
let slope = if header.scl_slope == 0.0 {
1.0
} else {
header.scl_slope
};
if slope != 1.0 || header.scl_inter != 0.0 {
let slope = A::from_f32(slope).unwrap();
let inter = A::from_f32(header.scl_inter).unwrap();
// TODO Use linear transformation like when reading. An scl_slope of 0.5 would turn all
// voxel values to 0 if we pass an ndarray of integers.
let slope = T::from_f32(slope).unwrap();
let inter = T::from_f32(header.scl_inter).unwrap();
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this part of the function should be reiterated (either now or in a separate PR), since it does not exactly resemble the current logic employed when reading transformed data. The data::element module provides multiple linear transformation strategies for each valid element data type to prevent precision loss. In this particular code, an scl_slope of 0.5 (haven't even seen real data with this, but it's a valid one nevertheless) would turn all voxel values to 0 if we pass an ndarray of integers. We might need a bit more code than just reusing this module though, since the transformation strategies currently live in immaterial types.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ouch. True. I'll check if I can copy your design.

for arr_data in data.axis_iter(Axis(0)) {
write_slice(arr_data.sub(inter).div(slope))?;
write_slice(writer, arr_data.sub(inter).div(slope))?;
}
} else {
for arr_data in data.axis_iter(Axis(0)) {
// We need to own the data because of the into_shape() for `c` ordering.
write_slice(arr_data.to_owned())?;
}
write_slices(writer, data)?;
}
Ok(())
}

#[cfg(test)]
pub mod tests {
extern crate tempfile;

use super::*;
use std::ops::{Add, Mul};
use std::path::PathBuf;

use self::tempfile::tempdir;
use ndarray::{Array, Array2, Ix2, IxDyn, ShapeBuilder};

use {header::MAGIC_CODE_NIP1, object::NiftiObject, InMemNiftiObject, IntoNdArray};

fn get_random_path(ext: &str) -> PathBuf {
let dir = tempdir().unwrap();
let mut path = dir.into_path();
if !ext.is_empty() {
path.push(ext);
}
path
}

pub fn generate_nifti_header(
dim: [u16; 8],
scl_slope: f32,
scl_inter: f32,
datatype: i16,
) -> NiftiHeader {
NiftiHeader {
dim,
datatype,
bitpix: (mem::size_of::<f32>() * 8) as i16,
magic: *MAGIC_CODE_NIP1,
scl_slope,
scl_inter,
..NiftiHeader::default()
}
}

fn read_2d_image<P>(path: P) -> Array2<f32>
where
P: AsRef<Path>,
{
let nifti_object = InMemNiftiObject::from_file(path).expect("Nifti file is unreadable.");
let volume = nifti_object.into_volume();
let dyn_data = volume.into_ndarray().unwrap();
dyn_data.into_dimensionality::<Ix2>().unwrap()
}

fn test_write_read(arr: &Array<f32, IxDyn>, path: &str) {
let path = get_random_path(path);
let mut dim = [1; 8];
dim[0] = arr.ndim() as u16;
for (i, s) in arr.shape().iter().enumerate() {
dim[i + 1] = *s as u16;
fn write_slices<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
S: Data<Elem = A>,
A: Clone + PodTransmutable,
D: Dimension + RemoveAxis,
W: WriteBytesExt,
{
let mut iter = data.axis_iter(Axis(0));
if let Some(arr_data) = iter.next() {
// Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
let mut slice = arr_data.to_owned();
write_slice(writer, slice.view())?;
for arr_data in iter {
slice.assign(&arr_data);
write_slice(writer, slice.view())?;
}
let header = generate_nifti_header(dim, 1.0, 0.0, 16);
write_nifti(&path, &arr, Some(&header)).unwrap();

let read_nifti = read_2d_image(path);
assert!(read_nifti.all_close(&arr, 1e-10));
}

fn f_order_array() -> Array<f32, IxDyn> {
let dim = vec![4, 4];
let vec = (0..16).map(|x| x as f32).collect();
Array::from_shape_vec(IxDyn(&dim).f(), vec).unwrap()
}

fn c_order_array() -> Array<f32, IxDyn> {
let dim = vec![4, 4];
let vec = (0..16).map(|x| x as f32).collect();
Array::from_shape_vec(IxDyn(&dim), vec).unwrap()
}

#[test]
fn test_fortran_writing() {
// Test .nii
let arr = f_order_array();
test_write_read(&arr, "test.nii");
let mut arr = f_order_array();
arr.invert_axis(Axis(1));
test_write_read(&arr, "test_non_contiguous.nii");

// Test .nii.gz
let arr = f_order_array();
test_write_read(&arr, "test.nii.gz");
let mut arr = f_order_array();
arr.invert_axis(Axis(1));
test_write_read(&arr, "test_non_contiguous.nii.gz");
}

#[test]
fn test_c_writing() {
// Test .nii
let arr = c_order_array();
test_write_read(&arr, "test.nii");

let mut arr = c_order_array();
arr.invert_axis(Axis(1));
test_write_read(&arr, "test_non_contiguous.nii");

// Test .nii.gz
let arr = c_order_array();
test_write_read(&arr, "test.nii.gz");

let mut arr = c_order_array();
arr.invert_axis(Axis(1));
test_write_read(&arr, "test_non_contiguous.nii.gz");
}
Ok(())
}

#[test]
fn test_header_slope_inter() {
let arr = f_order_array();
let slope = 2.2;
let inter = 101.1;

let path = get_random_path("test_slope_inter.nii");
let mut dim = [1; 8];
dim[0] = arr.ndim() as u16;
for (i, s) in arr.shape().iter().enumerate() {
dim[i + 1] = *s as u16;
}
let header = generate_nifti_header(dim, slope, inter, 16);
let transformed_data = arr.mul(slope).add(inter);
write_nifti(&path, &transformed_data, Some(&header)).unwrap();

let read_nifti = read_2d_image(path);
assert!(read_nifti.all_close(&transformed_data, 1e-10));
}
fn write_slice<A, S, D, W>(writer: &mut W, data: ArrayBase<S, D>) -> Result<()>
where
S: Data<Elem = A>,
A: Clone + PodTransmutable,
D: Dimension,
W: WriteBytesExt,
{
let len = data.len();
let arr_data = data.into_shape(len).unwrap();
let slice = arr_data.as_slice().unwrap();
writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
Ok(())
}
Loading