torchio.utils.compress() - Code Metrics - Inspection of "Bump version: 0.17.19 → 0.17.20" - fepegar/torchio - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Push — master ( 68ba61...357cca )

by Fernando

created 2020-08-05 10:13 UTC

torchio.utils.compress() A

↳ Parent: torchio.utils

Complexity

Conditions

Size

Total Lines	4
Code Lines	4

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	3
eloc	4
nop	2
dl	0
loc	4
rs	10
c	0
b	0
f	0

import ast
import gzip
import shutil
import tempfile
from pathlib import Path
from typing import Union, Iterable, Tuple, Any, Optional, List

import torch
import numpy as np
import nibabel as nib
import SimpleITK as sitk
from tqdm import trange
from .torchio import (
    INTENSITY,
    TypeData,
    TypeNumber,
    TypePath,
    REPO_URL,
)


FLIP_XY = np.diag((-1, -1, 1))  # used to switch between LPS and RAS


def to_tuple(
        value: Union[TypeNumber, Iterable[TypeNumber]],
        length: int = 1,
        ) -> Tuple[TypeNumber, ...]:
    """
    to_tuple(1, length=1) -> (1,)
    to_tuple(1, length=3) -> (1, 1, 1)

    If value is an iterable, n is ignored and tuple(value) is returned
    to_tuple((1,), length=1) -> (1,)
    to_tuple((1, 2), length=1) -> (1, 2)
    to_tuple([1, 2], length=3) -> (1, 2)
    """
    try:
        iter(value)
        value = tuple(value)
    except TypeError:
        value = length * (value,)
    return value


def get_stem(path: TypePath) -> str:
    """
    '/home/user/image.nii.gz' -> 'image'
    """
    path = Path(path)
    return path.name.split('.')[0]


def create_dummy_dataset(
        num_images: int,
        size_range: Tuple[int, int],
        directory: Optional[TypePath] = None,
        suffix: str = '.nii.gz',
        force: bool = False,
        verbose: bool = False,
        ):
    from .data import ScalarImage, LabelMap, Subject
    output_dir = tempfile.gettempdir() if directory is None else directory
    output_dir = Path(output_dir)
    images_dir = output_dir / 'dummy_images'
    labels_dir = output_dir / 'dummy_labels'

    if force:
        shutil.rmtree(images_dir)
        shutil.rmtree(labels_dir)

    subjects: List[Subject] = []
    if images_dir.is_dir():
        for i in trange(num_images):
            image_path = images_dir / f'image_{i}{suffix}'
            label_path = labels_dir / f'label_{i}{suffix}'
            subject = Subject(
                one_modality=ScalarImage(image_path),
                segmentation=LabelMap(label_path),
            )
            subjects.append(subject)
    else:
        images_dir.mkdir(exist_ok=True, parents=True)
        labels_dir.mkdir(exist_ok=True, parents=True)
        if verbose:
            print('Creating dummy dataset...')
            iterable = trange(num_images)
        else:
            iterable = range(num_images)
        for i in iterable:
            shape = np.random.randint(*size_range, size=3)
            affine = np.eye(4)
            image = np.random.rand(*shape)
            label = np.ones_like(image)
            label[image < 0.33] = 0
            label[image > 0.66] = 2
            image *= 255

            image_path = images_dir / f'image_{i}{suffix}'
            nii = nib.Nifti1Image(image.astype(np.uint8), affine)
            nii.to_filename(str(image_path))

            label_path = labels_dir / f'label_{i}{suffix}'
            nii = nib.Nifti1Image(label.astype(np.uint8), affine)
            nii.to_filename(str(label_path))

            subject = Subject(
                one_modality=ScalarImage(image_path),
                segmentation=LabelMap(label_path),
            )
            subjects.append(subject)
    return subjects


def apply_transform_to_file(
        input_path: TypePath,
        transform,  # : Transform seems to create a circular import (TODO)
        output_path: TypePath,
        type: str = INTENSITY,
        verbose: bool = False,
        ):
    from . import Image, ImagesDataset, Subject
    subject = Subject(image=Image(input_path, type=type))
    transformed = transform(subject)
    transformed.image.save(output_path)
    if verbose and transformed.history:
        print(transformed.history[0])

def guess_type(string: str) -> Any:
    # Adapted from
    # https://www.reddit.com/r/learnpython/comments/4599hl/module_to_guess_type_from_a_string/czw3f5s
    string = string.replace(' ', '')
    try:
        value = ast.literal_eval(string)
    except ValueError:
        result_type = str
    else:
        result_type = type(value)
    if result_type in (list, tuple):
        string = string[1:-1]  # remove brackets
        split = string.split(',')
        list_result = [guess_type(n) for n in split]
        value = tuple(list_result) if result_type is tuple else list_result
        return value
    try:
        value = result_type(string)
    except TypeError:
        value = None
    return value


def get_rotation_and_spacing_from_affine(
        affine: np.ndarray,
        ) -> Tuple[np.ndarray, np.ndarray]:
    # From https://github.com/nipy/nibabel/blob/master/nibabel/orientations.py
    rotation_zoom = affine[:3, :3]
    spacing = np.sqrt(np.sum(rotation_zoom * rotation_zoom, axis=0))
    rotation = rotation_zoom / spacing
    return rotation, spacing


def nib_to_sitk(
        data: TypeData,
        affine: TypeData,
        squeeze: bool = False,
        force_3d: bool = False,
        force_4d: bool = False,
        ) -> sitk.Image:
    """Create a SimpleITK image from a tensor and a 4x4 affine matrix.

    Args:
        data: PyTorch tensor or NumPy array
        affine: # TODO
    """
    if data.ndim != 4:
        raise ValueError(f'Input must be 4D, but has shape {tuple(data.shape)}')
    # Possibilities
    # (1, 1, h, w)
    # (c, 1, h, w)
    # (1, d, h, w)
    # (c, d, h, w)
    array = np.asarray(data)
    affine = np.asarray(affine).astype(np.float64)

    is_multichannel = array.shape[0] > 1 and not force_4d
    is_2d = array.shape[1] == 1 and not force_3d
    if is_2d:
        array = array[:, 0, :, :]
    if not is_multichannel and not force_4d:
        array = array[0]
    array = array.transpose()  # (W, H, D, C) or (W, H, D)
    image = sitk.GetImageFromArray(array, isVector=is_multichannel)

    rotation, spacing = get_rotation_and_spacing_from_affine(affine)
    origin = np.dot(FLIP_XY, affine[:3, 3])
    direction = np.dot(FLIP_XY, rotation)
    if is_2d:  # ignore first dimension if 2D (1, 1, H, W)
        direction = direction[1:3, 1:3]
    image.SetOrigin(origin)  # should I add a 4th value if force_4d?
    image.SetSpacing(spacing)
    image.SetDirection(direction.flatten())
    if data.ndim == 4:
        assert image.GetNumberOfComponentsPerPixel() == data.shape[0]
    num_spatial_dims = 2 if is_2d else 3
    assert image.GetSize() == data.shape[-num_spatial_dims:]
    return image


def sitk_to_nib(
        image: sitk.Image,
        keepdim: bool = False,
        ) -> Tuple[np.ndarray, np.ndarray]:
    data = sitk.GetArrayFromImage(image).transpose()
    num_components = image.GetNumberOfComponentsPerPixel()
    if num_components == 1:
        data = data[np.newaxis]  # add channels dimension
    input_spatial_dims = image.GetDimension()
    if not keepdim:
        data = ensure_4d(data, False, num_spatial_dims=input_spatial_dims)
    assert data.shape[0] == num_components
    assert data.shape[-input_spatial_dims:] == image.GetSize()
    spacing = np.array(image.GetSpacing())
    direction = np.array(image.GetDirection())
    origin = image.GetOrigin()
    if len(direction) == 9:
        rotation = direction.reshape(3, 3)
    elif len(direction) == 4:  # ignore first dimension if 2D (1, 1, H, W)
        rotation_2d = direction.reshape(2, 2)
        rotation = np.eye(3)
        rotation[1:3, 1:3] = rotation_2d
        spacing = 1, *spacing
        origin = 0, *origin
    rotation = np.dot(FLIP_XY, rotation)

    rotation_zoom = rotation * spacing
    translation = np.dot(FLIP_XY, origin)
    affine = np.eye(4)
    affine[:3, :3] = rotation_zoom
    affine[:3, 3] = translation
    return data, affine


def ensure_4d(
        tensor: TypeData,
        channels_last: bool,
        num_spatial_dims=None,
        ) -> TypeData:
    """[summary] # TODO

    Args:
        tensor: [description].
        channels_last: If ``True``, last dimension of the input represents
            channels.
        num_spatial_dims: [description].

    Raises:
        ValueError: [description]
    """
    # I wish named tensors were properly supported in PyTorch
    num_dimensions = tensor.ndim
    if num_dimensions == 5:  # hope (X, X, X, 1, X)
        if tensor.shape[-1] == 1:
            tensor = tensor[..., 0, :]
    if num_dimensions == 4:  # assume 3D multichannel
        if channels_last:  # (D, H, W, C)
            tensor = tensor.permute(3, 0, 1, 2)  # (C, D, H, W)
    elif num_dimensions == 2:  # assume 2D monochannel (H, W)
        tensor = tensor[np.newaxis, np.newaxis]  # (1, 1, H, W)
    elif num_dimensions == 3:  # 2D multichannel or 3D monochannel?
        if num_spatial_dims == 2:
            if channels_last:  # (H, W, C)
                tensor = tensor.permute(2, 0, 1)  # (C, H, W)
            tensor = tensor[:, np.newaxis]  # (C, 1, H, W)
        elif num_spatial_dims == 3:  # (D, H, W)
            tensor = tensor[np.newaxis]  # (1, D, H, W)
        else:  # try to guess
            shape = tensor.shape
            maybe_rgb = 3 in (shape[0], shape[-1])
            if maybe_rgb:
                if shape[-1] == 3:  # (H, W, 3)
                    tensor = tensor.permute(2, 0, 1)  # (3, H, W)
                tensor = tensor[:, np.newaxis]  # (3, 1, H, W)
            else:  # (D, H, W)
                tensor = tensor[np.newaxis]  # (1, D, H, W)
    else:
        message = (
            f'{num_dimensions}D images not supported yet. Please create an'
            f' issue in {REPO_URL} if you would like support for them'
        )
        raise ValueError(message)
    assert tensor.ndim == 4
    return tensor


def get_torchio_cache_dir():
    return Path('~/.cache/torchio').expanduser()


def round_up(value: float) -> float:
    """Round half towards infinity.

    Args:
        value: The value to round.

    Example:

        >>> round(2.5)
        2
        >>> round(3.5)
        4
        >>> round_up(2.5)
        3
        >>> round_up(3.5)
        4

    """
    return np.floor(value + 0.5)


def compress(input_path, output_path):
    with open(input_path, 'rb') as f_in:
        with gzip.open(output_path, 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)


1			import ast
2			import gzip
3			import shutil
4			import tempfile
5			from pathlib import Path
6			from typing import Union, Iterable, Tuple, Any, Optional, List
7
8			import torch
9			import numpy as np
10			import nibabel as nib
11			import SimpleITK as sitk
12			from tqdm import trange
13			from .torchio import (
14			INTENSITY,
15			TypeData,
16			TypeNumber,
17			TypePath,
18			REPO_URL,
19			)
20
21
22			FLIP_XY = np.diag((-1, -1, 1)) # used to switch between LPS and RAS
23
24
25			def to_tuple(
26			value: Union[TypeNumber, Iterable[TypeNumber]],
27			length: int = 1,
28			) -> Tuple[TypeNumber, ...]:
29			"""
30			to_tuple(1, length=1) -> (1,)
31			to_tuple(1, length=3) -> (1, 1, 1)
32
33			If value is an iterable, n is ignored and tuple(value) is returned
34			to_tuple((1,), length=1) -> (1,)
35			to_tuple((1, 2), length=1) -> (1, 2)
36			to_tuple([1, 2], length=3) -> (1, 2)
37			"""
38			try:
39			iter(value)
40			value = tuple(value)
41			except TypeError:
42			value = length * (value,)
43			return value
44
45
46			def get_stem(path: TypePath) -> str:
47			"""
48			'/home/user/image.nii.gz' -> 'image'
49			"""
50			path = Path(path)
51			return path.name.split('.')[0]
52
53
54			def create_dummy_dataset(
55			num_images: int,
56			size_range: Tuple[int, int],
57			directory: Optional[TypePath] = None,
58			suffix: str = '.nii.gz',
59			force: bool = False,
60			verbose: bool = False,
61			):
62			from .data import ScalarImage, LabelMap, Subject
63			output_dir = tempfile.gettempdir() if directory is None else directory
64			output_dir = Path(output_dir)
65			images_dir = output_dir / 'dummy_images'
66			labels_dir = output_dir / 'dummy_labels'
67
68			if force:
69			shutil.rmtree(images_dir)
70			shutil.rmtree(labels_dir)
71
72			subjects: List[Subject] = []
73			if images_dir.is_dir():
74			for i in trange(num_images):
75			image_path = images_dir / f'image_{i}{suffix}'
76			label_path = labels_dir / f'label_{i}{suffix}'
77			subject = Subject(
78			one_modality=ScalarImage(image_path),
79			segmentation=LabelMap(label_path),
80			)
81			subjects.append(subject)
82			else:
83			images_dir.mkdir(exist_ok=True, parents=True)
84			labels_dir.mkdir(exist_ok=True, parents=True)
85			if verbose:
86			print('Creating dummy dataset...')
87			iterable = trange(num_images)
88			else:
89			iterable = range(num_images)
90			for i in iterable:
91			shape = np.random.randint(*size_range, size=3)
92			affine = np.eye(4)
93			image = np.random.rand(*shape)
94			label = np.ones_like(image)
95			label[image < 0.33] = 0
96			label[image > 0.66] = 2
97			image *= 255
98
99			image_path = images_dir / f'image_{i}{suffix}'
100			nii = nib.Nifti1Image(image.astype(np.uint8), affine)
101			nii.to_filename(str(image_path))
102
103			label_path = labels_dir / f'label_{i}{suffix}'
104			nii = nib.Nifti1Image(label.astype(np.uint8), affine)
105			nii.to_filename(str(label_path))
106
107			subject = Subject(
108			one_modality=ScalarImage(image_path),
109			segmentation=LabelMap(label_path),
110			)
111			subjects.append(subject)
112			return subjects
113
114
115			def apply_transform_to_file(
116			input_path: TypePath,
117			transform, # : Transform seems to create a circular import (TODO)
118			output_path: TypePath,
119			type: str = INTENSITY,
120			verbose: bool = False,
121			):
122			from . import Image, ImagesDataset, Subject
123			subject = Subject(image=Image(input_path, type=type))
124			transformed = transform(subject)
125			transformed.image.save(output_path)
126			if verbose and transformed.history:
127			print(transformed.history[0])
128
129			def guess_type(string: str) -> Any:
130			# Adapted from
131			# https://www.reddit.com/r/learnpython/comments/4599hl/module_to_guess_type_from_a_string/czw3f5s
132			string = string.replace(' ', '')
133			try:
134			value = ast.literal_eval(string)
135			except ValueError:
136			result_type = str
137			else:
138			result_type = type(value)
139			if result_type in (list, tuple):
140			string = string[1:-1] # remove brackets
141			split = string.split(',')
142			list_result = [guess_type(n) for n in split]
143			value = tuple(list_result) if result_type is tuple else list_result
144			return value
145			try:
146			value = result_type(string)
147			except TypeError:
148			value = None
149			return value
150
151
152			def get_rotation_and_spacing_from_affine(
153			affine: np.ndarray,
154			) -> Tuple[np.ndarray, np.ndarray]:
155			# From https://github.com/nipy/nibabel/blob/master/nibabel/orientations.py
156			rotation_zoom = affine[:3, :3]
157			spacing = np.sqrt(np.sum(rotation_zoom * rotation_zoom, axis=0))
158			rotation = rotation_zoom / spacing
159			return rotation, spacing
160
161
162			def nib_to_sitk(
163			data: TypeData,
164			affine: TypeData,
165			squeeze: bool = False,
166			force_3d: bool = False,
167			force_4d: bool = False,
168			) -> sitk.Image:
169			"""Create a SimpleITK image from a tensor and a 4x4 affine matrix.
170
171			Args:
172			data: PyTorch tensor or NumPy array
173			affine: # TODO
174			"""
175			if data.ndim != 4:
176			raise ValueError(f'Input must be 4D, but has shape {tuple(data.shape)}')
177			# Possibilities
178			# (1, 1, h, w)
179			# (c, 1, h, w)
180			# (1, d, h, w)
181			# (c, d, h, w)
182			array = np.asarray(data)
183			affine = np.asarray(affine).astype(np.float64)
184
185			is_multichannel = array.shape[0] > 1 and not force_4d
186			is_2d = array.shape[1] == 1 and not force_3d
187			if is_2d:
188			array = array[:, 0, :, :]
189			if not is_multichannel and not force_4d:
190			array = array[0]
191			array = array.transpose() # (W, H, D, C) or (W, H, D)
192			image = sitk.GetImageFromArray(array, isVector=is_multichannel)
193
194			rotation, spacing = get_rotation_and_spacing_from_affine(affine)
195			origin = np.dot(FLIP_XY, affine[:3, 3])
196			direction = np.dot(FLIP_XY, rotation)
197			if is_2d: # ignore first dimension if 2D (1, 1, H, W)
198			direction = direction[1:3, 1:3]
199			image.SetOrigin(origin) # should I add a 4th value if force_4d?
200			image.SetSpacing(spacing)
201			image.SetDirection(direction.flatten())
202			if data.ndim == 4:
203			assert image.GetNumberOfComponentsPerPixel() == data.shape[0]
204			num_spatial_dims = 2 if is_2d else 3
205			assert image.GetSize() == data.shape[-num_spatial_dims:]
206			return image
207
208
209			def sitk_to_nib(
210			image: sitk.Image,
211			keepdim: bool = False,
212			) -> Tuple[np.ndarray, np.ndarray]:
213			data = sitk.GetArrayFromImage(image).transpose()
214			num_components = image.GetNumberOfComponentsPerPixel()
215			if num_components == 1:
216			data = data[np.newaxis] # add channels dimension
217			input_spatial_dims = image.GetDimension()
218			if not keepdim:
219			data = ensure_4d(data, False, num_spatial_dims=input_spatial_dims)
220			assert data.shape[0] == num_components
221			assert data.shape[-input_spatial_dims:] == image.GetSize()
222			spacing = np.array(image.GetSpacing())
223			direction = np.array(image.GetDirection())
224			origin = image.GetOrigin()
225			if len(direction) == 9:
226			rotation = direction.reshape(3, 3)
227			elif len(direction) == 4: # ignore first dimension if 2D (1, 1, H, W)
228			rotation_2d = direction.reshape(2, 2)
229			rotation = np.eye(3)
230			rotation[1:3, 1:3] = rotation_2d
231			spacing = 1, *spacing
232			origin = 0, *origin
233			rotation = np.dot(FLIP_XY, rotation)
			0 ignored issues – show introduced 2020-05-26 09:22 UTC by Report Bug Copy Issue Report The variable `rotation` does not seem to be defined for all execution paths. Loading history...
234			rotation_zoom = rotation * spacing
235			translation = np.dot(FLIP_XY, origin)
236			affine = np.eye(4)
237			affine[:3, :3] = rotation_zoom
238			affine[:3, 3] = translation
239			return data, affine
240
241
242			def ensure_4d(
243			tensor: TypeData,
244			channels_last: bool,
245			num_spatial_dims=None,
246			) -> TypeData:
247			"""[summary] # TODO
248
249			Args:
250			tensor: [description].
251			channels_last: If ``True``, last dimension of the input represents
252			channels.
253			num_spatial_dims: [description].
254
255			Raises:
256			ValueError: [description]
257			"""
258			# I wish named tensors were properly supported in PyTorch
259			num_dimensions = tensor.ndim
260			if num_dimensions == 5: # hope (X, X, X, 1, X)
261			if tensor.shape[-1] == 1:
262			tensor = tensor[..., 0, :]
263			if num_dimensions == 4: # assume 3D multichannel
264			if channels_last: # (D, H, W, C)
265			tensor = tensor.permute(3, 0, 1, 2) # (C, D, H, W)
266			elif num_dimensions == 2: # assume 2D monochannel (H, W)
267			tensor = tensor[np.newaxis, np.newaxis] # (1, 1, H, W)
268			elif num_dimensions == 3: # 2D multichannel or 3D monochannel?
269			if num_spatial_dims == 2:
270			if channels_last: # (H, W, C)
271			tensor = tensor.permute(2, 0, 1) # (C, H, W)
272			tensor = tensor[:, np.newaxis] # (C, 1, H, W)
273			elif num_spatial_dims == 3: # (D, H, W)
274			tensor = tensor[np.newaxis] # (1, D, H, W)
275			else: # try to guess
276			shape = tensor.shape
277			maybe_rgb = 3 in (shape[0], shape[-1])
278			if maybe_rgb:
279			if shape[-1] == 3: # (H, W, 3)
280			tensor = tensor.permute(2, 0, 1) # (3, H, W)
281			tensor = tensor[:, np.newaxis] # (3, 1, H, W)
282			else: # (D, H, W)
283			tensor = tensor[np.newaxis] # (1, D, H, W)
284			else:
285			message = (
286			f'{num_dimensions}D images not supported yet. Please create an'
287			f' issue in {REPO_URL} if you would like support for them'
288			)
289			raise ValueError(message)
290			assert tensor.ndim == 4
291			return tensor
292
293
294			def get_torchio_cache_dir():
295			return Path('~/.cache/torchio').expanduser()
296
297
298			def round_up(value: float) -> float:
299			"""Round half towards infinity.
300
301			Args:
302			value: The value to round.
303
304			Example:
305
306			>>> round(2.5)
307			2
308			>>> round(3.5)
309			4
310			>>> round_up(2.5)
311			3
312			>>> round_up(3.5)
313			4
314
315			"""
316			return np.floor(value + 0.5)
317
318
319			def compress(input_path, output_path):
320			with open(input_path, 'rb') as f_in:
321			with gzip.open(output_path, 'wb') as f_out:
322			shutil.copyfileobj(f_in, f_out)
323

fepegar / torchio

Push — master ( 68ba61...357cca )

torchio.utils.compress() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like