From f6e0248b458e0fb7f2049ba01e9e002906c7614e Mon Sep 17 00:00:00 2001 From: ancestor-mithril Date: Wed, 3 Apr 2024 14:07:27 +0300 Subject: [PATCH] Improving image io * reducing the number of allocations when reading images and writing segmentations * segmentation are casted to np.uint8 only if they are not already np.uint8 * images are stacked directly into a float32 array, without additional casts --- nnunetv2/imageio/natural_image_reader_writer.py | 4 ++-- nnunetv2/imageio/nibabel_reader_writer.py | 8 +++----- nnunetv2/imageio/simpleitk_reader_writer.py | 5 ++--- nnunetv2/imageio/tif_reader_writer.py | 6 +++--- 4 files changed, 10 insertions(+), 13 deletions(-) diff --git a/nnunetv2/imageio/natural_image_reader_writer.py b/nnunetv2/imageio/natural_image_reader_writer.py index 11946c3ca..7ce42e8a4 100644 --- a/nnunetv2/imageio/natural_image_reader_writer.py +++ b/nnunetv2/imageio/natural_image_reader_writer.py @@ -56,13 +56,13 @@ def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[ print('Image files:') print(image_fnames) raise RuntimeError() - return np.vstack(images).astype(np.float32), {'spacing': (999, 1, 1)} + return np.vstack(images, dtype=np.float32, casting='unsafe'), {'spacing': (999, 1, 1)} def read_seg(self, seg_fname: str) -> Tuple[np.ndarray, dict]: return self.read_images((seg_fname, )) def write_seg(self, seg: np.ndarray, output_fname: str, properties: dict) -> None: - io.imsave(output_fname, seg[0].astype(np.uint8), check_contrast=False) + io.imsave(output_fname, seg[0].astype(np.uint8, copy=False), check_contrast=False) if __name__ == '__main__': diff --git a/nnunetv2/imageio/nibabel_reader_writer.py b/nnunetv2/imageio/nibabel_reader_writer.py index 8faafb709..5ffa5488a 100644 --- a/nnunetv2/imageio/nibabel_reader_writer.py +++ b/nnunetv2/imageio/nibabel_reader_writer.py @@ -78,14 +78,13 @@ def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[ print(image_fnames) raise RuntimeError() - stacked_images = np.vstack(images) dict = { 'nibabel_stuff': { 'original_affine': original_affines[0], }, 'spacing': spacings_for_nnunet[0] } - return stacked_images.astype(np.float32), dict + return np.vstack(images, dtype=np.float32, casting='unsafe'), dict def read_seg(self, seg_fname: str) -> Tuple[np.ndarray, dict]: return self.read_images((seg_fname, )) @@ -160,7 +159,6 @@ def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[ print(image_fnames) raise RuntimeError() - stacked_images = np.vstack(images) dict = { 'nibabel_stuff': { 'original_affine': original_affines[0], @@ -168,14 +166,14 @@ def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[ }, 'spacing': spacings_for_nnunet[0] } - return stacked_images.astype(np.float32), dict + return np.vstack(images, dtype=np.float32, casting='unsafe'), dict def read_seg(self, seg_fname: str) -> Tuple[np.ndarray, dict]: return self.read_images((seg_fname, )) def write_seg(self, seg: np.ndarray, output_fname: str, properties: dict) -> None: # revert transpose - seg = seg.transpose((2, 1, 0)).astype(np.uint8) + seg = seg.transpose((2, 1, 0)).astype(np.uint8, copy=False) seg_nib = nibabel.Nifti1Image(seg, affine=properties['nibabel_stuff']['reoriented_affine']) seg_nib_reoriented = seg_nib.as_reoriented(io_orientation(properties['nibabel_stuff']['original_affine'])) diff --git a/nnunetv2/imageio/simpleitk_reader_writer.py b/nnunetv2/imageio/simpleitk_reader_writer.py index 6a9afc24f..0f386e5da 100644 --- a/nnunetv2/imageio/simpleitk_reader_writer.py +++ b/nnunetv2/imageio/simpleitk_reader_writer.py @@ -97,7 +97,6 @@ def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[ print(image_fnames) raise RuntimeError() - stacked_images = np.vstack(images) dict = { 'sitk_stuff': { # this saves the sitk geometry information. This part is NOT used by nnU-Net! @@ -109,7 +108,7 @@ def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[ # are returned x,y,z but spacing is returned z,y,x. Duh. 'spacing': spacings_for_nnunet[0] } - return stacked_images.astype(np.float32), dict + return np.vstack(images, dtype=np.float32, casting='unsafe'), dict def read_seg(self, seg_fname: str) -> Tuple[np.ndarray, dict]: return self.read_images((seg_fname, )) @@ -121,7 +120,7 @@ def write_seg(self, seg: np.ndarray, output_fname: str, properties: dict) -> Non if output_dimension == 2: seg = seg[0] - itk_image = sitk.GetImageFromArray(seg.astype(np.uint8)) + itk_image = sitk.GetImageFromArray(seg.astype(np.uint8, copy=False)) itk_image.SetSpacing(properties['sitk_stuff']['spacing']) itk_image.SetOrigin(properties['sitk_stuff']['origin']) itk_image.SetDirection(properties['sitk_stuff']['direction']) diff --git a/nnunetv2/imageio/tif_reader_writer.py b/nnunetv2/imageio/tif_reader_writer.py index 19ad882a3..b1623495a 100644 --- a/nnunetv2/imageio/tif_reader_writer.py +++ b/nnunetv2/imageio/tif_reader_writer.py @@ -66,11 +66,11 @@ def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[ print(image_fnames) raise RuntimeError() - return np.vstack(images).astype(np.float32), {'spacing': spacing} + return np.vstack(images, dtype=np.float32, casting='unsafe'), {'spacing': spacing} def write_seg(self, seg: np.ndarray, output_fname: str, properties: dict) -> None: # not ideal but I really have no clue how to set spacing/resolution information properly in tif files haha - tifffile.imwrite(output_fname, data=seg.astype(np.uint8), compression='zlib') + tifffile.imwrite(output_fname, data=seg.astype(np.uint8, copy=False), compression='zlib') file = os.path.basename(output_fname) out_dir = os.path.dirname(output_fname) ending = file.split('.')[-1] @@ -97,4 +97,4 @@ def read_seg(self, seg_fname: str) -> Tuple[np.ndarray, dict]: print(f'WARNING no spacing file found for segmentation {seg_fname}\nAssuming spacing (1, 1, 1).') spacing = (1, 1, 1) - return seg.astype(np.float32), {'spacing': spacing} \ No newline at end of file + return seg.astype(np.float32, copy=False), {'spacing': spacing}