torchio.utils.gen_seed() - Code Metrics - Inspection of "Add some deterministic transforms" - fepegar/torchio - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — master (#353)

by Fernando

created 2020-11-17 02:13 UTC

torchio.utils.gen_seed() A

↳ Parent: torchio.utils

Complexity

Conditions

Size

Total Lines	7
Code Lines	2

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
eloc	2
nop	0
dl	0
loc	7
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, Sequence

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: Union[TypePath, List[TypePath]]
        ) -> Union[str, List[str]]:
    """
    '/home/user/image.nii.gz' -> 'image'
    """
    def _get_stem(path_string):
        return Path(path_string).name.split('.')[0]
    if isinstance(path, (str, Path)):
        return _get_stem(path)
    return [_get_stem(p) for p in path]


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...')  # noqa: T001
            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
        output_path: TypePath,
        type: str = INTENSITY,  # noqa: A002
        verbose: bool = False,
        ):
    from . import Image, 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])  # noqa: T001


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."""
    if data.ndim != 4:
        raise ValueError(f'Input must be 4D, but has shape {tuple(data.shape)}')
    # Possibilities
    # (1, w, h, 1)
    # (c, w, h, 1)
    # (1, w, h, 1)
    # (c, w, h, d)
    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[3] == 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, W, H, 1)
        direction = direction[:2, :2]
    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[1: 1 + 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 input_spatial_dims == 2:
        data = data[..., np.newaxis]
    if not keepdim:
        data = ensure_4d(data, num_spatial_dims=input_spatial_dims)
    assert data.shape[0] == num_components
    assert data.shape[1: 1 + 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, W, H, 1)
        rotation_2d = direction.reshape(2, 2)
        rotation = np.eye(3)
        rotation[:2, :2] = rotation_2d
        spacing = *spacing, 1
        origin = *origin, 0
    else:
        raise RuntimeError(f'Direction not understood: {direction}')
    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, num_spatial_dims=None) -> TypeData:
    # I wish named tensors were properly supported in PyTorch
    num_dimensions = tensor.ndim
    if num_dimensions == 4:
        pass
    elif num_dimensions == 5:  # hope (X, X, X, 1, X)
        if tensor.shape[-1] == 1:
            tensor = tensor[..., 0, :]
    elif num_dimensions == 2:  # assume 2D monochannel (W, H)
        tensor = tensor[np.newaxis, ..., np.newaxis]  # (1, W, H, 1)
    elif num_dimensions == 3:  # 2D multichannel or 3D monochannel?
        if num_spatial_dims == 2:
            tensor = tensor[..., np.newaxis]  # (C, W, H, 1)
        elif num_spatial_dims == 3:  # (W, H, D)
            tensor = tensor[np.newaxis]  # (1, W, H, D)
        else:  # try to guess
            shape = tensor.shape
            maybe_rgb = 3 in (shape[0], shape[-1])
            if maybe_rgb:
                if shape[-1] == 3:  # (W, H, 3)
                    tensor = tensor.permute(2, 0, 1)  # (3, W, H)
                tensor = tensor[..., np.newaxis]  # (3, W, H, 1)
            else:  # (W, H, D)
                tensor = tensor[np.newaxis]  # (1, W, H, D)
    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) -> int:
    """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 int(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)


def check_sequence(sequence: Sequence, name: str):
    try:
        iter(sequence)
    except TypeError:
        message = f'"{name}" must be a sequence, not {type(name)}'
        raise TypeError(message)


def get_random_seed():
    """Generate a random seed.

    Returns:
        A random seed as an int
    """
    return torch.randint(0, 2**31, (1,)).item()


def get_major_sitk_version() -> int:
    # This attribute was added in version 2
    # https://github.com/SimpleITK/SimpleITK/pull/1171
    version = getattr(sitk, '__version__', None)
    major_version = 1 if version is None else 2
    return major_version


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

fepegar / torchio

Pull Request — master (#353)

torchio.utils.gen_seed() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like