File size: 3,128 Bytes
64fd08f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | import torchio as tio
import numpy as np
import torch
# Helper function for multiprocessing-safe casting
def _to_float32(tensor: torch.Tensor) -> torch.Tensor:
"""Cast TorchIO tensor to float32 (multiprocessing-safe)."""
return tensor.to(torch.float32)
class PadToMultiple(tio.Transform):
"""
Pads a 3D image so that each spatial dimension is a multiple of a given number.
"""
def __init__(self, multiple: int = 16, **kwargs):
super().__init__(**kwargs)
self.multiple = multiple
def apply_transform(self, subject: tio.Subject) -> tio.Subject:
image_keys = sorted([i for i in subject.keys() if i.startswith("image_")])
for image_key in image_keys:
image = subject[image_key]
shape = image.spatial_shape
target_shape = [int(np.ceil(dim / self.multiple)) * self.multiple for dim in shape]
pad_transform = tio.CropOrPad(
target_shape=tuple(target_shape), padding_mode=0
)
subject = pad_transform(subject)
return subject
def pad_fixed_size_collate_fn(batch):
"""
Pads each image in the batch to the same size and returns a tensor of images and a tensor of labels.
Assumes that the subject only contains one image and one label.
"""
# Find the maximum spatial dimensions in the batch
max_shape = np.max([item['image_1'][tio.DATA].shape[1:] for item in batch], axis=0)
# Make it a multiple of 16
target_shape = tuple([int(np.ceil(dim / 16)) * 16 for dim in max_shape])
pad_transform = tio.CropOrPad(
target_shape=target_shape, padding_mode=0
)
# Pad each image to the max shape
for c, item in enumerate(batch):
image_keys = sorted([i for i in item.keys() if i.startswith("image_")])
for image_key in image_keys:
image = item[image_key]
padded_image = pad_transform(image)
batch[c][image_key] = padded_image
return batch
def get_image_transforms(do_augmentation: bool = True):
"""
Get the image transforms for the training and validation sets.
"""
if do_augmentation:
train_transforms = tio.Compose([
tio.RandomFlip(axes=('LR', 'AP', 'IS'), flip_probability=0.5),
tio.RandomAffine(scales=(0.8, 1.2), degrees=45, translation=(15, 15, 15), p=0.3),
tio.RandomNoise(std=(0.01, 0.10), p=0.2),
# tio.RandomBiasField(coefficients=0.5, p=0.2),
# tio.RandomBlur(std=(0, 2), p=0.2),
tio.RandomGamma(log_gamma=(-0.3, 0.3), p=0.2),
tio.RandomSwap(patch_size=10, num_iterations=30, p=0.2)
])
else:
train_transforms = tio.Compose([
tio.RandomFlip(axes=('LR', 'AP', 'IS'), flip_probability=0.5),
tio.RandomAffine(scales=(0.8, 1.2), degrees=45, translation=(15, 15, 15), p=0.3),
])
# Add type transform to ensure image data loading for validation
val_transforms = tio.Compose([tio.Lambda(_to_float32)])
return train_transforms, val_transforms
|