diff --git a/examples/trk_len_check.rs b/examples/trk_len_check.rs new file mode 100644 index 0000000..2b709ac --- /dev/null +++ b/examples/trk_len_check.rs @@ -0,0 +1,80 @@ +use std::time::Instant; +use std::{fs::File, path::Path, str}; + +use docopt::Docopt; + +use trk_io::Reader; + +static USAGE: &'static str = " +Binary to test that reserving a number of points, scalars and properties is +possible and works 100% of the time. + +Usage: + trk_len_check [options] + trk_len_check (-h | --help) + trk_len_check (-v | --version) + +Options: + -h --help Show this screen. + -v --version Show version. +"; + +fn main() { + let version = String::from(env!("CARGO_PKG_VERSION")); + let args = Docopt::new(USAGE) + .and_then(|dopt| dopt.version(Some(version)).parse()) + .unwrap_or_else(|e| e.exit()); + let input = Path::new(args.get_str("")); + if !input.exists() { + panic!("Input trk '{:?}' doesn't exist.", input); + } + + let f = File::open(args.get_str("")).expect("Can't read trk file."); + let metadata = f.metadata().unwrap(); + let nb_bytes = metadata.len() as usize; + + let reader = Reader::new(args.get_str("")).expect("Read header"); + let header = &reader.header; + let nb_scalars = header.scalars_name.len(); + let nb_properties = header.properties_name.len(); + let nb_streamlines = header.nb_streamlines; + let header_size = header.raw_header().hdr_size as usize; + + let nb_f32s = (nb_bytes - header_size) / 4; + let guessed_nb_properties = nb_streamlines * nb_properties; + + // We remove the nb_points count and the properties, se we only have the points and scalars + let nb_f32s = nb_f32s - nb_streamlines - guessed_nb_properties; + let guessed_nb_points = nb_f32s / (3 + nb_scalars); + let guessed_nb_scalars = nb_f32s - (3 * guessed_nb_points); + + let mut real_nb_points = 0; + let mut real_nb_scalars = 0; + let mut real_nb_properties = 0; + for (streamline, scalars, properties) in reader.into_iter() { + real_nb_points += streamline.len(); + for s in &scalars { + real_nb_scalars += s.len(); + } + real_nb_properties += properties.len(); + } + + assert_eq!(guessed_nb_points, real_nb_points); + assert_eq!(guessed_nb_scalars, real_nb_scalars); + assert_eq!(guessed_nb_properties, real_nb_properties); + + let mut reader = Reader::new(args.get_str("")).expect("Read header"); + let now1 = Instant::now(); + let _ = reader.read_all(); + let elapsed1 = now1.elapsed(); + println!("{}s {}ms", elapsed1.as_secs(), elapsed1.subsec_millis()); + + let mut reader = Reader::new(args.get_str("")).expect("Read header"); + let now2 = Instant::now(); + let _ = reader.read_all_alloc(); + let elapsed2 = now2.elapsed(); + println!("{}s {}ms", elapsed2.as_secs(), elapsed2.subsec_millis()); + + let diff = elapsed1 - elapsed2; + println!("{}ms OR {}ns", diff.subsec_millis(), diff.subsec_nanos()) +} diff --git a/src/reader.rs b/src/reader.rs index 819d753..ac6c374 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -20,6 +20,7 @@ pub struct Reader { pub affine_to_rasmm: Affine, pub translation: Translation, + nb_bytes: usize, nb_scalars: usize, nb_properties: usize, nb_floats_per_point: usize, @@ -28,7 +29,10 @@ pub struct Reader { impl Reader { pub fn new>(path: P) -> Result { - let mut reader = BufReader::new(File::open(&path)?); + let f = File::open(&path)?; + let nb_bytes = f.metadata().unwrap().len() as usize; + + let mut reader = BufReader::new(f); let (header, endianness) = Header::read(&mut reader)?; let affine_to_rasmm = header.affine_to_rasmm; let translation = header.translation; @@ -42,6 +46,7 @@ impl Reader { header, affine_to_rasmm, translation, + nb_bytes, nb_scalars, nb_properties, nb_floats_per_point, @@ -66,7 +71,6 @@ impl Reader { } fn read_all_(&mut self) -> Tractogram { - // TODO Anything we can do to reerve? let mut lengths = Vec::new(); let mut v = Vec::with_capacity(300); let mut scalars = ArraySequence::with_capacity(300); @@ -81,6 +85,41 @@ impl Reader { Tractogram::new(Streamlines::new(lengths, v), scalars, properties) } + pub fn read_all_alloc(&mut self) -> Tractogram { + match self.endianness { + Endianness::Little => self.read_all_alloc_::(), + Endianness::Big => self.read_all_alloc_::(), + } + } + + fn read_all_alloc_(&mut self) -> Tractogram { + let nb_scalars = self.header.scalars_name.len(); + let nb_properties = self.header.properties_name.len(); + let nb_streamlines = self.header.nb_streamlines; + let header_size = self.header.raw_header().hdr_size as usize; + + // We remove the nb_points count and the properties from nb_f32s, se we only have the + // points and scalars + let guessed_nb_properties = nb_streamlines * nb_properties; + let nb_f32s = (self.nb_bytes - header_size) / 4; + let nb_f32s = nb_f32s - nb_streamlines - guessed_nb_properties; + let guessed_nb_points = nb_f32s / (3 + nb_scalars); + let guessed_nb_scalars = nb_f32s - (3 * guessed_nb_points); + + let mut lengths = Vec::with_capacity(nb_streamlines); + let mut v = Vec::with_capacity(guessed_nb_points); + let mut scalars = ArraySequence::with_capacity(guessed_nb_scalars); + let mut properties = ArraySequence::with_capacity(guessed_nb_properties); + while let Ok(nb_points) = self.reader.read_i32::() { + lengths.push(nb_points as usize); + self.read_streamline::(&mut v, &mut scalars, nb_points as usize); + self.read_properties_to_arr::(&mut properties); + } + + self.float_buffer = vec![]; + Tractogram::new(Streamlines::new(lengths, v), scalars, properties) + } + fn read_streamline( &mut self, points: &mut Points,