File size: 4,593 Bytes
5a93c05 | 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | import numpy as np
import os
import random
import torch
from scipy.ndimage.interpolation import zoom
from torch.utils.data import Dataset
from scipy import ndimage
from PIL import Image
from scipy.ndimage import map_coordinates, gaussian_filter
import cv2
def random_rot_flip(image, label):
k = np.random.randint(0, 4)
image = np.rot90(image, k, axes=(0,1))
label = np.rot90(label, k, axes=(0,1))
axis = np.random.randint(0, 2)
image = np.flip(image, axis=axis)
label = np.flip(label, axis=axis)
return image.copy(), label.copy()
def random_rotate(image, label):
angle = np.random.uniform(-25, 25)
rotated_channels = []
for c in range(image.shape[-1]):
rotated_c = ndimage.rotate(
image[..., c], angle, order=1,
reshape=False, mode='nearest'
)
rotated_channels.append(rotated_c)
image_rot = np.stack(rotated_channels, axis=-1)
label_rot = ndimage.rotate(
label, angle, order=0,
reshape=False, mode='nearest'
)
return image_rot, label_rot
def elastic_transform(image, label, alpha=1000, sigma=30):
shape = image.shape[:2]
random_state = np.random.RandomState(None)
dx = gaussian_filter(
(random_state.rand(*shape) * 2 - 1),
sigma, mode="constant"
) * alpha
dy = gaussian_filter(
(random_state.rand(*shape) * 2 - 1),
sigma, mode="constant"
) * alpha
x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij')
indices = np.reshape(x + dx, (-1, 1)), np.reshape(y + dy, (-1, 1))
dist_image = []
for c in range(image.shape[-1]):
channel = map_coordinates(
image[..., c], indices,
order=3, mode='reflect'
).reshape(shape)
dist_image.append(channel)
dist_image = np.stack(dist_image, axis=-1)
dist_label = map_coordinates(
label, indices,
order=0, mode='reflect'
).reshape(shape)
return dist_image, dist_label
class AdvancedMedicalAug(torch.nn.Module):
def __init__(self, aug_prob=0.8):
super().__init__()
self.aug_prob = aug_prob
def forward(self, image, label):
if isinstance(image, torch.Tensor):
image = image.numpy()
if isinstance(label, torch.Tensor):
label = label.numpy()
if random.random() < self.aug_prob:
if random.random() > 0.5:
image, label = random_rot_flip(image, label)
else:
image, label = random_rotate(image, label)
if random.random() < 0.3:
image, label = elastic_transform(image, label)
image = self.intensity_augment(image)
return image, label
def intensity_augment(self, image):
for c in range(image.shape[-1]):
image[..., c] = np.clip(
image[..., c] * random.uniform(0.7, 1.3),
0, 1
)
if random.random() < 0.2:
noise = np.random.normal(0, 0.05, image[..., c].shape)
image[..., c] = np.clip(image[..., c] + noise, 0, 1)
return image
class Synapse_dataset(Dataset):
def __init__(self, base_dir, list_dir, split, img_size, transform=None):
self.transform = transform
self.split = split
self.sample_list = open(os.path.join(list_dir, self.split+'.txt')).readlines()
self.data_dir = base_dir
self.img_size = img_size
def __len__(self):
return len(self.sample_list)
def __getitem__(self, idx):
slice_name = self.sample_list[idx].strip('\n')
data_path = os.path.join(self.data_dir, slice_name)
data = np.load(data_path)
image = data['image'].astype(np.float32)
if len(image.shape) == 2:
image = np.expand_dims(image, axis=-1)
if 'label' in data.files:
label = data['label'].astype(np.int32)
else:
label = np.zeros_like(image, dtype=np.int32)
if image.shape != self.img_size:
x, y, z = image.shape
image = zoom(image, (self.img_size[0] / x, self.img_size[1] / y, z), order=3)
label = zoom(label, (self.img_size[0] / x, self.img_size[1] / y), order=0)
if self.transform:
image, label = self.transform(image, label)
if self.split == 'train':
image, label = torch.from_numpy(image), torch.from_numpy(label)
sample = {'image': image, 'label': label, 'case_name': slice_name}
return sample
|