torchio.utils.create_dummy_dataset() - Code Metrics - Inspection of "Add support for 4D images" - fepegar/torchio - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — master (#246)

by Fernando

created 2020-08-02 22:04 UTC

torchio.utils.create_dummy_dataset() B

↳ Parent: torchio.utils

Complexity

Conditions

Size

Total Lines	59
Code Lines	49

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	7
eloc	49
nop	6
dl	0
loc	59
rs	7.269
c	0
b	0
f	0

How to fix Long Method

import ast
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,
    LABEL,
    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 Image, 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=Image(image_path, INTENSITY),
                segmentation=Image(label_path, LABEL),
            )
            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=Image(image_path, INTENSITY),
                segmentation=Image(label_path, LABEL),
            )
            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))
    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,
        ) -> 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.array(data)
    affine = np.array(affine).astype(np.float64)

    is_multichannel = array.shape[0] > 1
    is_2d = array.shape[1] == 1 and not force_3d
    if is_2d:
        array = array[:, 0, :, :]
    if not is_multichannel:
        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)
    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]:
    """[summary]

    Args:
        image (sitk.Image): [description]
        keepdim (bool, optional): [description]. Defaults to False.

    Returns:
        Tuple[np.ndarray, np.ndarray]: [description]
    """
    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
            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)


1			import ast
2			import shutil
3			import tempfile
4			from pathlib import Path
5			from typing import Union, Iterable, Tuple, Any, Optional, List
6
7			import torch
8			import numpy as np
9			import nibabel as nib
10			import SimpleITK as sitk
11			from tqdm import trange
12			from .torchio import (
13			INTENSITY,
14			LABEL,
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 Image, 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=Image(image_path, INTENSITY),
79			segmentation=Image(label_path, LABEL),
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=Image(image_path, INTENSITY),
109			segmentation=Image(label_path, LABEL),
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))
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			) -> sitk.Image:
168			"""Create a SimpleITK image from a tensor and a 4x4 affine matrix.
169
170			Args:
171			data: PyTorch tensor or NumPy array
172			affine: # TODO
173			"""
174			if data.ndim != 4:
175			raise ValueError(f'Input must be 4D, but has shape {tuple(data.shape)}')
176			# Possibilities
177			# (1, 1, h, w)
178			# (c, 1, h, w)
179			# (1, d, h, w)
180			# (c, d, h, w)
181			array = np.array(data)
182			affine = np.array(affine).astype(np.float64)
183
184			is_multichannel = array.shape[0] > 1
185			is_2d = array.shape[1] == 1 and not force_3d
186			if is_2d:
187			array = array[:, 0, :, :]
188			if not is_multichannel:
189			array = array[0]
190			array = array.transpose() # (W, H, D, C) or (W, H, D)
191			image = sitk.GetImageFromArray(array, isVector=is_multichannel)
192
193			rotation, spacing = get_rotation_and_spacing_from_affine(affine)
194			origin = np.dot(FLIP_XY, affine[:3, 3])
195			direction = np.dot(FLIP_XY, rotation)
196			if is_2d: # ignore first dimension if 2D (1, 1, H, W)
197			direction = direction[1:3, 1:3]
198			image.SetOrigin(origin)
199			image.SetSpacing(spacing)
200			image.SetDirection(direction.flatten())
201			if data.ndim == 4:
202			assert image.GetNumberOfComponentsPerPixel() == data.shape[0]
203			num_spatial_dims = 2 if is_2d else 3
204			assert image.GetSize() == data.shape[-num_spatial_dims:]
205			return image
206
207
208			def sitk_to_nib(
209			image: sitk.Image,
210			keepdim: bool = False,
211			) -> Tuple[np.ndarray, np.ndarray]:
212			"""[summary]
213
214			Args:
215			image (sitk.Image): [description]
216			keepdim (bool, optional): [description]. Defaults to False.
217
218			Returns:
219			Tuple[np.ndarray, np.ndarray]: [description]
220			"""
221			data = sitk.GetArrayFromImage(image).transpose()
222			num_components = image.GetNumberOfComponentsPerPixel()
223			if num_components == 1:
224			data = data[np.newaxis] # add channels dimension
225			input_spatial_dims = image.GetDimension()
226			if not keepdim:
227			data = ensure_4d(data, False, num_spatial_dims=input_spatial_dims)
228			assert data.shape[0] == num_components
229			assert data.shape[-input_spatial_dims:] == image.GetSize()
230			spacing = np.array(image.GetSpacing())
231			direction = np.array(image.GetDirection())
232			origin = image.GetOrigin()
233			if len(direction) == 9:
234			rotation = direction.reshape(3, 3)
235			elif len(direction) == 4: # ignore first dimension if 2D (1, 1, H, W)
236			rotation_2d = direction.reshape(2, 2)
237			rotation = np.eye(3)
238			rotation[1:3, 1:3] = rotation_2d
239			spacing = 1, *spacing
240			origin = 0, *origin
241			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...
242			rotation_zoom = rotation * spacing
243			translation = np.dot(FLIP_XY, origin)
244			affine = np.eye(4)
245			affine[:3, :3] = rotation_zoom
246			affine[:3, 3] = translation
247			return data, affine
248
249
250			def ensure_4d(
251			tensor: TypeData,
252			channels_last: bool,
253			num_spatial_dims=None,
254			) -> TypeData:
255			"""[summary] # TODO
256
257			Args:
258			tensor: [description].
259			channels_last: If ``True``, last dimension of the input represents
260			channels.
261			num_spatial_dims: [description].
262
263			Raises:
264			ValueError: [description]
265			"""
266			# I wish named tensors were properly supported in PyTorch
267			num_dimensions = tensor.ndim
268			if num_dimensions == 5: # hope (X, X, X, 1, X)
269			if tensor.shape[-1] == 1:
270			tensor = tensor[..., 0, :]
271			if num_dimensions == 4: # assume 3D multichannel
272			if channels_last: # (D, H, W, C)
273			tensor = tensor.permute(3, 0, 1, 2) # (C, D, H, W)
274			elif num_dimensions == 2: # assume 2D monochannel (H, W)
275			tensor = tensor[np.newaxis, np.newaxis] # (1, 1, H, W)
276			elif num_dimensions == 3: # 2D multichannel or 3D monochannel?
277			if num_spatial_dims == 2:
278			if channels_last: # (H, W, C)
279			tensor = tensor.permute(2, 0, 1) # (C, H, W)
280			tensor = tensor[:, np.newaxis] # (C, 1, H, W)
281			elif num_spatial_dims == 3: # (D, H, W)
282			tensor = tensor[np.newaxis] # (1, D, H, W)
283			else: # try to guess
284			shape = tensor.shape
285			maybe_rgb = 3 in shape
286			if maybe_rgb:
287			if shape[-1] == 3: # (H, W, 3)
288			tensor = tensor.permute(2, 0, 1) # (3, H, W)
289			tensor = tensor[:, np.newaxis] # (3, 1, H, W)
290			else: # (D, H, W)
291			tensor = tensor[np.newaxis] # (1, D, H, W)
292			else:
293			message = (
294			f'{num_dimensions}D images not supported yet. Please create an'
295			f' issue in {REPO_URL} if you would like support for them'
296			)
297			raise ValueError(message)
298			assert tensor.ndim == 4
299			return tensor
300
301
302			def get_torchio_cache_dir():
303			return Path('~/.cache/torchio').expanduser()
304
305
306			def round_up(value: float) -> float:
307			"""Round half towards infinity.
308
309			Args:
310			value: The value to round.
311
312			Example:
313
314			>>> round(2.5)
315			2
316			>>> round(3.5)
317			4
318			>>> round_up(2.5)
319			3
320			>>> round_up(3.5)
321			4
322
323			"""
324			return np.floor(value + 0.5)
325

fepegar / torchio

Pull Request — master (#246)

torchio.utils.create_dummy_dataset() B

Complexity

Size

Duplication

Importance

How to fix Long Method

Long Method

Duplication Side-by-Side

Filter issues like