def augmentation(self, flow, normalize=True):
"""
Take an existing flow of data and augment it with consistent random transformations uniformly accross all
frames in the video. Also supports normalization /in [0,1]
:param flow: the iterator of video data
:param normalize: do you want to scale the data to [0,1] using a linear map
"""
image_datagen = ImageDataGenerator(**self.data_gen_args)
for video in flow:
# for every frame in this video generate the same transformation
# and yield it all back in sequence order
trans = image_datagen.get_random_transform(video[2].shape)
augmentedVideo = np.zeros(video[2].shape)
for i in range(video[2].shape[0]):
augmentedVideo[i] = image_datagen.apply_transform(
video[2][i], trans)
# now is a good time to transform the video onto 0-1
# we need to do this to get convergence when we train i.e. homogenise features
if normalize:
augmentedVideo[i] = augmentedVideo[i] / 255
yield video[:-1] + (augmentedVideo, )
def augment(movie, mask):
#MOVIE is a 3d array that represents each slice of the movie stacked on top of each other
# transform every horizontal slice
datagen_xy = ImageDataGenerator(zoom_range=[.75, 1.1],
rotation_range=90,
shear_range=10,
horizontal_flip=True,
vertical_flip=True)
xy_transform = datagen_xy.get_random_transform(movie.shape[1:3])
for i in range(movie.shape[0]):
movie[i, :, :, :] = datagen_xy.apply_transform(movie[i, :, :, :],
xy_transform)
#transform the mask to match
mask = datagen_xy.apply_transform(mask, xy_transform)
return movie, mask
def augment(movie, mask): #MOVIE is a 3d array that represents each slice of the movie stacked on top of each other #first, transform every vertical slice # datagen_zy = ImageDataGenerator(zoom_range = .15) # z_transform = datagen_zy.get_random_transform(movie.shape[0:2]) # z_transform["zx"] = 0 # for i in range(movie.shape[2]): # movie[:, :, i, :] = datagen_zy.apply_transform(movie[:, :, i, :], z_transform) #now transform every horizontal slice datagen_xy = ImageDataGenerator(zoom_range = [.6, 1], rotation_range = 90, shear_range = 12, horizontal_flip = True, vertical_flip = True) xy_transform = datagen_xy.get_random_transform(movie.shape[1:3]) for i in range(movie.shape[0]): movie[i, :, :, :] = datagen_xy.apply_transform(movie[i, :, :, :], xy_transform) #transform the mask to match mask = datagen_xy.apply_transform(mask, xy_transform) return movie, mask
def transform(inputs, outputs, ntimes=8, args=None):
datagen = ImageDataGenerator(**args)
input_gen = []
output_gen = []
for i in range(ntimes):
for j in range(len(inputs)):
inp = inputs[j]
out = outputs[j]
trans = datagen.get_random_transform(inp.shape)
inp = datagen.apply_transform(inp, trans)
out = datagen.apply_transform(out, trans)
input_gen.append(inp)
output_gen.append(out)
input_gen = np.array(input_gen)
output_gen = np.array(output_gen)
return input_gen, output_gen
class DataLoaderCamus:
def __init__(self, dataset_path, input_name, target_name, condition_name,
img_res, target_rescale, input_rescale, condition_rescale, train_ratio, valid_ratio,
labels, augment):
self.dataset_path = dataset_path
self.img_res = tuple(img_res)
self.target_rescale = target_rescale
self.input_rescale = input_rescale
self.condition_rescale = condition_rescale
self.input_name = input_name
self.target_name = target_name
self.condition_name = condition_name
self.augment = augment
patients = sorted(glob(os.path.join(self.dataset_path, 'training', '*')))
random.Random(RANDOM_SEED).shuffle(patients)
num = len(patients)
num_train = int(num * train_ratio)
num_valid = int(num_train * valid_ratio)
self.valid_patients = patients[:num_valid]
self.train_patients = patients[num_valid:num_train]
self.test_patients = patients[num_train:]
if train_ratio == 1.0:
self.test_patients = glob(os.path.join(self.dataset_path, 'testing', '*'))
print('#train:', len(self.train_patients))
print('#valid:', len(self.valid_patients))
print('#test:', len(self.test_patients))
all_labels = {0, 1, 2, 3}
self.not_labels = all_labels - set(labels)
data_gen_args = dict(rotation_range=augment['AUG_ROTATION_RANGE_DEGREES'],
width_shift_range=augment['AUG_WIDTH_SHIFT_RANGE_RATIO'],
height_shift_range=augment['AUG_HEIGHT_SHIFT_RANGE_RATIO'],
shear_range=augment['AUG_SHEAR_RANGE_ANGLE'],
zoom_range=augment['AUG_ZOOM_RANGE_RATIO'],
fill_mode='constant',
cval=0.,
data_format='channels_last')
self.datagen = ImageDataGenerator(**data_gen_args)
def read_mhd(self, img_path, is_gt):
if not os.path.exists(img_path):
return np.zeros(self.img_res + (1,))
img = io.imread(img_path, plugin='simpleitk').squeeze()
img = np.array(Image.fromarray(img).resize(self.img_res))
img = np.expand_dims(img, axis=2)
if is_gt:
for not_l in self.not_labels:
img[img == not_l] = 0
return img
def _get_paths(self, stage):
if stage == 'train':
return self.train_patients
elif stage == 'valid':
return self.valid_patients
elif stage == 'test':
return self.test_patients
@background(max_prefetch=NUM_PREFETCH)
def get_random_batch(self, batch_size=1, stage='train'):
paths = self._get_paths(stage)
num = len(paths)
num_batches = num // batch_size
for i in range(num_batches):
batch_paths = np.random.choice(paths, size=batch_size)
target_imgs, condition_imgs, input_imgs, weight_imgs = self._get_batch(batch_paths, stage)
target_imgs = target_imgs * self.target_rescale
input_imgs = input_imgs * self.input_rescale
condition_imgs = condition_imgs * self.condition_rescale
yield target_imgs, condition_imgs, input_imgs, weight_imgs
def get_iterative_batch(self, batch_size=1, stage='test'):
paths = self._get_paths(stage)
num = len(paths)
num_batches = num // batch_size
start_idx = 0
for i in range(num_batches):
batch_paths = paths[start_idx:start_idx + batch_size]
target_imgs, condition_imgs, input_imgs, weight_imgs = self._get_batch(batch_paths, stage)
target_imgs = target_imgs * self.target_rescale
input_imgs = input_imgs * self.input_rescale
condition_imgs = condition_imgs * self.condition_rescale
start_idx += batch_size
yield target_imgs, condition_imgs, input_imgs, weight_imgs
def _get_batch(self, paths_batch, stage):
target_imgs = []
input_imgs = []
condition_imgs = []
weight_maps = []
for path in paths_batch:
transform = self.datagen.get_random_transform(img_shape=self.img_res)
head, patient_id = os.path.split(path)
target_path = os.path.join(path, '{}_{}.mhd'.format(patient_id, self.target_name))
condition_path = os.path.join(path, '{}_{}.mhd'.format(patient_id, self.condition_name))
input_path = os.path.join(path, '{}_{}.mhd'.format(patient_id, self.input_name))
input_img = self.read_mhd(input_path, '_gt' in self.input_name)
if self.augment['AUG_INPUT']:
input_img = self.datagen.apply_transform(input_img, transform)
input_imgs.append(input_img)
target_img = self.read_mhd(target_path, '_gt' in self.target_name)
condition_img = self.read_mhd(condition_path, 1)
if self.augment['AUG_TARGET']:
if not self.augment['AUG_SAME_FOR_BOTH']:
transform = self.datagen.get_random_transform(img_shape=self.img_res)
target_img = self.datagen.apply_transform(target_img, transform)
condition_img = self.datagen.apply_transform(condition_img, transform)
target_imgs.append(target_img)
condition_imgs.append(condition_img)
weight_map_condition = self.get_weight_map(condition_img)
weight_maps.append(weight_map_condition)
return np.array(target_imgs), np.array(condition_imgs), np.array(input_imgs), np.array(weight_maps)
def get_weight_map(self, mask):
# let the y axis have higher variance
gauss_var = [[self.img_res[0] * 60, 0], [0, self.img_res[1] * 30]]
x, y = mask[:, :, 0].nonzero()
center = [x.mean(), y.mean()]
from scipy.stats import multivariate_normal
gauss = multivariate_normal.pdf(np.mgrid[
0:self.img_res[1],
0:self.img_res[0]].reshape(2, -1).transpose(),
mean=center,
cov=gauss_var)
gauss /= gauss.max()
gauss = gauss.reshape((self.img_res[1], self.img_res[0], 1))
# set the gauss value of the main target part to 1
gauss[mask > 0] = 1
return gauss
def mj_getNegLAEOpair(negsamples, videoname, timepos, winlen, meanSample=[0], imgsize=(64,64)):
'''
Gets just one pair of negative samples
:param negsamples:
:param videoname:
:param timepos:
:param winlen:
:param imgsize: (rows, cols) of output crop
:return: output images are already normalized (x/255)
'''
foo = 0
nvids = len(negsamples["videoname"])
vix = -1
for vix_ in range(0,nvids):
if negsamples["videoname"][vix_] == videoname:
vix = vix_
if vix >= 0:
lTr = negsamples["tracks"][vix]
if timepos >= len(lTr):
return None, None
ltrx = lTr[timepos]
ntracks = len(ltrx)
if ntracks < 2:
return None, None
# TODO: parametrize these values, currently, random
if ntracks > 2:
rnp = np.random.permutation(range(0, ntracks))
t1 = rnp[0]
t2 = rnp[1]
else:
t1 = 0
t2 = 1
cropsvid = negsamples["crops"][vix]
geomvid = negsamples["geom"][vix]
pair = np.zeros((winlen, imgsize[0], 2*imgsize[1],3)) # Allocate memory for the temporal sequence
if len(ltrx[t1]) < winlen or len(ltrx[t2]) < winlen: # Just in case
return None, None
# Define an image transformation
from keras.preprocessing.image import ImageDataGenerator
img_gen = ImageDataGenerator(width_shift_range=[-2, 0, 2], height_shift_range=[-2, 0, 2],
brightness_range=[0.95, 1.05], channel_shift_range=0.05,
zoom_range=0.015, horizontal_flip=True)
transf = img_gen.get_random_transform(cropsvid[0][0].shape)
transf["flip_horizontal"] = False
G = 0
for tix in range(timepos,timepos+winlen):
ix1 = ltrx[t1][tix-timepos]
ix2 = ltrx[t2][tix-timepos]
if tix >= len(geomvid):
return None, None
# Check which one is on the left
if geomvid[tix][ix1][0] > geomvid[tix][ix2][0]:
ix1, ix2 = ix2, ix1 # Swap
crop1 = copy.deepcopy(cropsvid[tix][ix1])
crop1 = img_gen.apply_transform(crop1, transf)
crop2 = copy.deepcopy(cropsvid[tix][ix2])
crop2 = img_gen.apply_transform(crop2, transf)
geo1 = geomvid[tix][ix1]
geo2 = geomvid[tix][ix2]
dx = geo2[0]-geo1[0]
dy = geo2[1]-geo1[1]
rscale = geo1[2] / geo2[2]
crop1 = crop1/255.0
crop2 = crop2/255.0
if crop1.shape[0] != imgsize[0]:
crop1 = cv2.resize(crop1, imgsize)
crop2 = cv2.resize(crop2, imgsize)
if type(meanSample) == np.ndarray and meanSample.shape[0] == crop1.shape[0]:
crop1 -= meanSample
crop2 -= meanSample
p = np.concatenate((crop1,crop2), axis=1)
pair[tix - timepos,] = p
#cv2.imshow("Pair", p/255)
#cv2.waitKey(-1)
G = G + np.array([dx, dy, rscale])
imgsize_ = negsamples["imsize"][vix]
G = G / winlen
G[0] = G[0] / imgsize_[1]
G[1] = G[1] / imgsize_[0]
if winlen == 1:
pair = np.squeeze(pair)
return pair, G
else:
return None, None
class TextImageGenerator:
def __init__(self, img_dirpath, labels_path, img_w, img_h,
batch_size, downsample_factor, idxs, training=True, max_text_len=9, n_eraser=5):
self.img_h = img_h
self.img_w = img_w
self.batch_size = batch_size
self.max_text_len = max_text_len
self.idxs = idxs
self.downsample_factor = downsample_factor
self.img_dirpath = img_dirpath # image dir path
self.labels= json.load(open(labels_path))
self.img_dir = os.listdir(self.img_dirpath) # images list
self.img_dir = [self.img_dir[idx] for idx in self.idxs]
self.n = len(self.img_dir) # number of images
self.indexes = list(range(self.n))
self.cur_index = 0
self.imgs = np.zeros((self.n, self.img_h, self.img_w, 3), dtype=np.float32)
self.training = training
self.n_eraser = n_eraser
self.random_eraser = get_random_eraser(s_l=0.004, s_h=0.005, r_1=0.01, r_2=1/0.01, v_l=-128, v_h=128)
self.texts = []
image_datagen_args = {
'shear_range': 0.1,
'zoom_range': 0.01,
'width_shift_range': 0.001,
'height_shift_range': 0.1,
'rotation_range': 1,
'horizontal_flip': False,
'vertical_flip': False
}
self.image_datagen = ImageDataGenerator(**image_datagen_args)
def build_data(self):
print(self.n, " Image Loading start...")
for i, img_file in enumerate(self.img_dir):
img = cv2.imread(os.path.join(self.img_dirpath, img_file), cv2.IMREAD_GRAYSCALE)
# Add random padding
# img = random_padding(img, max_width_height_ratio=20, min_width_height_ratio=10, chanels=1)
# Resize & black white
img = cv2.resize(img, (self.img_w, self.img_h))
(thresh, img) = cv2.threshold(img, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)
img = img.astype(np.float32)
img = preprocess_input(img)
self.imgs[i] = img
self.texts.append(self.labels[img_file])
print("Image Loading finish...")
def next_sample(self):
self.cur_index += 1
if self.cur_index >= self.n:
self.cur_index = 0
random.shuffle(self.indexes)
return self.imgs[self.indexes[self.cur_index]], self.texts[self.indexes[self.cur_index]]
def next_batch(self):
while True:
X_data = np.zeros([self.batch_size, self.img_w, self.img_h, 3], dtype=np.float32) # (bs, 128, 64, 1)
Y_data = np.zeros([self.batch_size, self.max_text_len], dtype=np.float32) # (bs, 9)
input_length = np.ones((self.batch_size, 1), dtype=np.float32) * (self.img_w // self.downsample_factor - 2) # (bs, 1)
label_length = np.zeros((self.batch_size, 1), dtype=np.float32) # (bs, 1)
for i in range(self.batch_size):
img, text = self.next_sample()
if self.training:
params = self.image_datagen.get_random_transform(img.shape)
img = self.image_datagen.apply_transform(img, params)
for _ in range(self.n_eraser):
img = self.random_eraser(img)
img = img.transpose((1, 0, 2))
# random eraser if training
X_data[i] = img
Y_data[i,:len(text)] = text_to_labels(text)
label_length[i] = len(text)
inputs = {
# 'the_inputs': X_data, # (bs, 128, 64, 1)
'input_1': X_data, # (bs, 128, 64, 1)
'the_labels': Y_data, # (bs, 8)
'input_length': input_length, # (bs, 1)
'label_length': label_length # (bs, 1)
}
outputs = {'ctc': np.zeros([self.batch_size])} # (bs, 1)
yield (inputs, outputs)
class DataGenerator(Sequence):
def __init__(self,
video_directory,
batch_size,
is_training,
rgb_data_only=False,
spectrogram_dir=None):
if not rgb_data_only and not spectrogram_dir:
raise ValueError(
"spectrogram_dir is required if rgb_data_only is set to False."
)
self.video_filepaths = self._get_filepaths(video_directory)
self.batch_size = batch_size
self.is_training = is_training
self.rgb_data_only = rgb_data_only
self.spectrogram_dir = spectrogram_dir
if self.is_training:
np.random.shuffle(self.video_filepaths)
self.data_augmentor = ImageDataGenerator(
horizontal_flip=True,
brightness_range=[0.7, 1.3],
zoom_range=0.3)
else:
# self.labels are used as y_true when calculating model metrics
self.labels = [
config.RGB_CLASS_NAME_TO_IDX[self._get_label_name(filepath)]
for filepath in self.video_filepaths
]
def __getitem__(self, batch_num):
n_filepaths = len(self.video_filepaths)
idx_start = batch_num * self.batch_size
idx_end = min((batch_num + 1) * self.batch_size, n_filepaths)
filepaths_for_batch = self.video_filepaths[idx_start:idx_end]
X, y = self._get_batch(filepaths_for_batch)
return X, y
def __len__(self):
return int(np.ceil(len(self.video_filepaths) / self.batch_size))
def on_epoch_end(self):
if self.is_training:
np.random.shuffle(self.video_filepaths)
def _apply_data_augmentation(self, frames):
frames_augmented = []
img_shape = (config.RGB_FRAME_HEIGHT, config.RGB_FRAME_WIDTH,
config.CHANNELS)
transform_params = self.data_augmentor.get_random_transform(img_shape)
for frame in frames:
frame_augmented = self.data_augmentor.apply_transform(
frame, transform_params)
frame_augmented = frame_augmented / 255.0
frames_augmented.append(frame_augmented)
return np.array(frames_augmented)
def _class_directory_is_class(self, class_directory):
return class_directory.split("/")[-1] in config.RGB_CLASS_NAME_TO_IDX
def _crop_frame(self, frame, crop_location):
"""
crop_location=0.0 -> far left (landscape frame) or far top (portrait frame)
crop_location=0.5 -> center crop
crop_location==1.0 -> far right (landscape frame) or far bottom (portrait frame)
"""
h, w = frame.shape[:2]
length = min(h, w)
y_margin = h - length
y0 = int(y_margin * crop_location)
y1 = y0 + length
x_margin = w - length
x0 = int(x_margin * crop_location)
x1 = x0 + length
frame_cropped = frame[y0:y1, x0:x1]
frame_resized = cv2.resize(
frame_cropped, (config.RGB_FRAME_HEIGHT, config.RGB_FRAME_WIDTH))
return frame_resized
def _crop_frames(self, frames, center_crop=True):
"""If center_crop is False, the crop location will be random."""
cropped_frames = []
crop_location = 0.5 if center_crop else np.random.random_sample()
for frame in frames:
cropped_frame = self._crop_frame(frame, crop_location)
cropped_frames.append(cropped_frame)
return np.array(cropped_frames)
def _extract_frames_from_video(self, video_filepath):
video = cv2.VideoCapture(video_filepath)
if not video.isOpened():
raise FileNotFoundError(
"The input video path you provided is invalid.")
frames = []
while video.isOpened():
grabbed, frame_bgr = video.read()
if not grabbed:
break
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
frames.append(frame_rgb)
video.release()
return frames
def _get_batch(self, batch_video_filepaths):
batch_frames = [self._get_frames(fp) for fp in batch_video_filepaths]
batch_frames = np.array(batch_frames)
batch_labels = [
config.RGB_CLASS_NAME_TO_IDX[self._get_label_name(fp)]
for fp in batch_video_filepaths
]
batch_labels = np.array(batch_labels)
if self.rgb_data_only:
X, y = batch_frames, batch_labels
else:
batch_spectrograms = [
self._get_spectrogram(fp) for fp in batch_video_filepaths
]
batch_spectrograms = np.array(batch_spectrograms)
X, y = [batch_frames, batch_spectrograms], batch_labels
return X, y
def _get_frames(self, video_filepath):
frames = self._extract_frames_from_video(video_filepath)
frames = self._pad_frames_list(frames)
frames = self._subsample_frames(frames)
frames = self._crop_frames(frames, center_crop=not self.is_training)
if self.is_training:
frames = self._apply_data_augmentation(frames)
else:
frames = frames / 255.0
return frames
def _get_filepaths(self, directory):
filepaths = []
class_pathname = os.path.join(directory, "*")
class_dirs = glob.glob(class_pathname)
for class_dir in class_dirs:
if not self._class_directory_is_class(class_dir):
continue
pathname = os.path.join(class_dir, "*")
filepaths += glob.glob(pathname)
return filepaths
def _get_label_name(self, filepath):
return filepath.split("/")[-2]
def _get_spectrogram(self, video_filepath):
video_label = self._get_label_name(video_filepath)
audio_label = config.VIDEO_TO_AUDIO_LABEL_MAPPING[video_label]
video_filename = video_filepath.split('/')[-1].split('.')[0]
spectrogram_filepath = f'{self.spectrogram_dir}/{audio_label}/{video_label}_{video_filename}.jpg'
spectrogram_img_bgr = cv2.imread(spectrogram_filepath)
spectrogram_img = cv2.cvtColor(spectrogram_img_bgr, cv2.COLOR_BGR2RGB)
spectrogram_img = spectrogram_img / 255.0
spectrogram_img = cv2.resize(
spectrogram_img,
(config.SPECTROGRAM_HEIGHT, config.SPECTROGRAM_WIDTH))
return spectrogram_img
def _pad_frames_list(self, frames):
"""If the length of frames list is less than RGB_N_FRAMES,
it will be padded with blank frames (RGB -> 000).
"""
if len(frames) < config.RGB_N_FRAMES:
n_pad_frames = config.RGB_N_FRAMES - len(frames)
for _ in range(n_pad_frames):
blank_frame = np.zeros(
(config.RGB_FRAME_HEIGHT, config.RGB_FRAME_WIDTH,
config.CHANNELS))
frames.append(blank_frame)
return frames
def _subsample_frames(self, video_clip_frames):
"""Frames are subsampled uniformly. i.e. A fixed number of frames are
subsampled from video_clip_frames with equal distance from each other.
"""
subsampled_frames = []
current_ix = 0
step_size = len(video_clip_frames) / float(config.RGB_N_FRAMES)
for _ in range(config.RGB_N_FRAMES):
frame = video_clip_frames[int(current_ix)]
subsampled_frames.append(frame)
current_ix += step_size
return np.array(subsampled_frames)
class TextImageGenerator:
def __init__(self, img_dirpath, labels_path, img_w, img_h,
batch_size, downsample_factor, idxs, training=True, max_text_len=9, n_eraser=5):
self.img_h = img_h
self.img_w = img_w
self.batch_size = batch_size
self.max_text_len = max_text_len
self.idxs = idxs
self.downsample_factor = downsample_factor
self.img_dirpath = img_dirpath # image dir path
self.labels= json.load(open(labels_path)) if labels_path != None else None
self.img_dir = sorted(os.listdir(self.img_dirpath)) # images list
random.shuffle(self.img_dir)
if self.idxs is not None:
self.img_dir = [self.img_dir[idx] for idx in self.idxs]
self.n = len(self.img_dir) # number of images
self.indexes = list(range(self.n))
self.cur_index = 0
self.imgs = np.ones((self.n, self.img_h, self.img_w, 3), dtype=np.float16)
self.training = training
self.n_eraser = n_eraser
self.random_eraser = get_random_eraser(s_l=0.004, s_h=0.005, r_1=0.01, r_2=1/0.01, v_l=-128, v_h=128)
self.texts = []
image_datagen_args = {
'shear_range': 0.1,
'zoom_range': 0.01,
'width_shift_range': 0.001,
'height_shift_range': 0.1,
'rotation_range': 1,
'horizontal_flip': False,
'vertical_flip': False
}
self.image_datagen = ImageDataGenerator(**image_datagen_args)
def build_data(self):
print(self.n, " Image Loading start... ", self.img_dirpath)
for i, img_file in enumerate(self.img_dir):
img = image.load_img(self.img_dirpath + img_file, target_size=SIZE[::-1], interpolation='bicubic')
img = image.img_to_array(img)
img = preprocess_input(img)
self.imgs[i] = img
if self.labels != None:
self.texts.append(self.labels[img_file][:MAX_LEN])
else:
#valid mode
self.texts.append('')
print("Image Loading finish...")
def next_sample(self):
self.cur_index += 1
if self.cur_index >= self.n:
self.cur_index = 0
random.shuffle(self.indexes)
return self.imgs[self.indexes[self.cur_index]].astype(np.float32), self.texts[self.indexes[self.cur_index]]
def next_batch(self):
while True:
X_data = np.zeros([self.batch_size, self.img_w, self.img_h, 3], dtype=np.float32) # (bs, 128, 64, 1)
Y_data = np.zeros([self.batch_size, self.max_text_len], dtype=np.float32) # (bs, 9)
input_length = np.ones((self.batch_size, 1), dtype=np.float32) * (self.img_w // self.downsample_factor - 2) # (bs, 1)
label_length = np.zeros((self.batch_size, 1), dtype=np.float32) # (bs, 1)
for i in range(self.batch_size):
img, text = self.next_sample()
if self.training:
params = self.image_datagen.get_random_transform(img.shape)
img = self.image_datagen.apply_transform(img, params)
if randint(0, 1) == 1:
for _ in range(self.n_eraser):
img = self.random_eraser(img)
img = elastic_transform(img, 10, 2, 0.1)
img = img.transpose((1, 0, 2))
# random eraser if training
X_data[i] = img
Y_data[i,:len(text)] = text_to_labels(text)
label_length[i] = len(text)
inputs = {
'the_inputs': X_data, # (bs, 128, 64, 1)
'the_labels': Y_data, # (bs, 8)
'input_length': input_length, # (bs, 1)
'label_length': label_length # (bs, 1)
}
outputs = {'ctc': np.zeros([self.batch_size])} # (bs, 1)
yield (inputs, outputs)
class DataLoaderCamus:
def __init__(self, dataset_path, input_name, target_name, img_res,
target_rescale, input_rescale, train_ratio, valid_ratio,
labels, augment, equalize_lv_length):
self.dataset_path = dataset_path
self.img_res = tuple(img_res)
self.target_rescale = target_rescale
self.input_rescale = input_rescale
self.input_name = input_name
self.target_name = target_name
self.augment = augment
self.equalize_lv_length = equalize_lv_length
patients = sorted(
glob(os.path.join(self.dataset_path, 'training', '*')))
random.Random(RANDOM_SEED).shuffle(patients)
num = len(patients)
num_train = int(num * train_ratio)
valid_num = int(num_train * valid_ratio)
self.valid_patients = patients[:valid_num]
self.train_patients = patients[valid_num:num_train]
self.test_patients = patients[num_train:]
print('#train:', len(self.train_patients))
print('#valid:', len(self.valid_patients))
print('#test:', len(self.test_patients))
print('Consistency check - First valid sample:',
self.valid_patients[0])
print('Consistency check - First test sample:', self.test_patients[0])
all_labels = {0, 1, 2, 3}
self.not_labels = all_labels - set(labels)
data_gen_args = dict(
rotation_range=augment['AUG_ROTATION_RANGE_DEGREES'],
width_shift_range=augment['AUG_WIDTH_SHIFT_RANGE_RATIO'],
height_shift_range=augment['AUG_HEIGHT_SHIFT_RANGE_RATIO'],
shear_range=augment['AUG_SHEAR_RANGE_ANGLE'],
zoom_range=augment['AUG_ZOOM_RANGE_RATIO'],
fill_mode='constant',
cval=0.,
data_format='channels_last')
self.datagen = ImageDataGenerator(**data_gen_args)
def read_mhd(self, img_path, is_gt):
if not os.path.exists(img_path):
return np.zeros(self.img_res + (1, ))
img = io.imread(img_path, plugin='simpleitk').squeeze()
img = np.array(Image.fromarray(img).resize(self.img_res))
img = np.expand_dims(img, axis=2)
if is_gt:
for not_l in self.not_labels:
img[img == not_l] = 0
return img
def _get_paths(self, stage):
if stage == 'train':
return self.train_patients
elif stage == 'valid':
return self.valid_patients
elif stage == 'test':
return self.test_patients
@background(max_prefetch=NUM_PREFETCH)
def get_random_batch(self, batch_size=1, stage='train'):
paths = self._get_paths(stage)
num = len(paths)
num_batches = num // batch_size
for i in range(num_batches):
batch_paths = np.random.choice(paths, size=batch_size)
target_imgs, target_imgs_gt, input_imgs, input_imgs_gt = self._get_batch(
batch_paths, stage)
target_imgs = target_imgs * self.target_rescale
input_imgs = input_imgs * self.input_rescale
yield target_imgs, target_imgs_gt, input_imgs, input_imgs_gt
def get_iterative_batch(self, batch_size=1, stage='test'):
paths = self._get_paths(stage)
num = len(paths)
num_batches = num // batch_size
start_idx = 0
for i in range(num_batches):
batch_paths = paths[start_idx:start_idx + batch_size]
target_imgs, target_imgs_gt, input_imgs, input_imgs_gt, = self._get_batch(
batch_paths, stage)
target_imgs = target_imgs * self.target_rescale
input_imgs = input_imgs * self.input_rescale
start_idx += batch_size
yield target_imgs, target_imgs_gt, input_imgs, input_imgs_gt
def _get_batch(self, paths_batch, stage):
target_imgs = []
source_imgs = []
target_imgs_gt = []
source_gt_imgs = []
for path in paths_batch:
transform = self.datagen.get_random_transform(
img_shape=self.img_res)
head, patient_id = os.path.split(path)
target_path = os.path.join(
path, '{}_{}.mhd'.format(patient_id, self.target_name))
target_gt_path = os.path.join(
path, '{}_{}.mhd'.format(patient_id, self.target_name + '_gt'))
source_path = os.path.join(
path, '{}_{}.mhd'.format(patient_id, self.input_name))
source_gt_path = os.path.join(
path, '{}_{}.mhd'.format(patient_id, self.input_name + '_gt'))
# get source
source_img = self.read_mhd(source_path, '_gt' in self.input_name)
source_gt_img = self.read_mhd(source_gt_path, 1)
if stage == 'train':
source_img = self.datagen.apply_transform(
source_img, transform)
source_gt_img = self.datagen.apply_transform(
source_gt_img, transform)
# get target
target_img = self.read_mhd(target_path, '_gt' in self.target_name)
target_gt_img = self.read_mhd(target_gt_path, 1)
if self.augment['AUG_TARGET'] and stage == 'train':
if not self.augment['AUG_SAME_FOR_BOTH']:
transform = self.datagen.get_random_transform(
img_shape=self.img_res)
target_img = self.datagen.apply_transform(
target_img, transform)
target_gt_img = self.datagen.apply_transform(
target_gt_img, transform)
# equalize LV height of source to target
if self.equalize_lv_length:
source_img, source_gt_img = self.equalize_lv(
target_gt_img, source_img, source_gt_img)
# add to list
source_imgs.append(source_img)
source_gt_imgs.append(source_gt_img)
target_imgs.append(target_img)
target_imgs_gt.append(target_gt_img)
np.array(source_imgs)
return np.array(target_imgs), np.array(target_imgs_gt), np.array(
source_imgs), np.array(source_gt_imgs)
def equalize_lv(self, target_gt_img, source_img, source_gt_img):
def resize_img(img, ratio):
img = cv2.resize(img, (0, 0), fx=ratio, fy=ratio)
img = match_image_size(img, self.img_res)
img = np.expand_dims(img, -1)
assert img.shape[0] == target_gt_img.shape[0] and img.shape[
1] == target_gt_img.shape[1]
return img
# calculate ratio to resize
source_lv_length, _, _ = get_LV_lenght(source_gt_img.squeeze(), True)
target_lv_length, _, _ = get_LV_lenght(target_gt_img.squeeze(), True)
ratio = target_lv_length / source_lv_length
# resize source image and gt image
source_img = resize_img(source_img, ratio)
source_gt_img = resize_img(source_gt_img, ratio)
return source_img, source_gt_img
class generator3da(Sequence):
def __init__(self,
list_IDs,
image_path,
mask_path,
to_fit=True,
batch_size=32,
patch_size=8,
dim=(512, 512),
dimy=(512, 512),
n_channels=1,
n_classes=10,
shuffle=True,
data_gen_args=None):
"""Initialization
:param list_IDs: list of all 'label' ids to use in the generator
:param image_path: path to images location
:param mask_path: path to masks location
:param to_fit: True to return X and y, False to return X only
:param batch_size: batch size at each iteration
:param dim: tuple indicating image dimension
:param n_channels: number of image channels
:param n_classes: number of output masks
:param shuffle: True to shuffle label indexes after every epoch
"""
self.patch_size = patch_size
self.number_of_patches = 0
self.list_IDs = list_IDs
self.image_path = image_path
self.mask_path = mask_path
self.to_fit = to_fit
self.batch_size = batch_size
self.dim = dim
self.dimy = dimy
self.n_channels = n_channels
self.n_classes = n_classes
self.shuffle = shuffle
self.on_epoch_end()
self.n = 0
self.max = self.__len__()
slices = os.listdir(os.path.join(self.image_path, self.list_IDs[0]))
self.number_of_patches = int(np.floor((len(slices) / self.patch_size)))
self.patientIDs = list_IDs
self.list_IDs = []
if data_gen_args != None:
self.trans = ImageDataGenerator(**data_gen_args)
temp = []
count = 0
for i, ID in enumerate(self.patientIDs):
slices = os.listdir(os.path.join(self.image_path, ID))
while count < self.number_of_patches:
patch = slices[(count * self.patch_size):((count + 1) *
self.patch_size)]
self.list_IDs.append([ID, patch])
count += 1
count = 0
self.indexes = np.arange(len(self.list_IDs))
def _load_dicom_image(self, image_path):
"""Load grayscale image
:param image_path: path to image to load
:return: loaded image
"""
img = load_dicom(image_path)
img = img / np.amax(img)
#self.polar(img)
return img
def _load_grayscale_image_VTK(self, image_path):
"""Load grayscale image
:param image_path: path to image to load
:return: loaded image
"""
img = vtk.vtkPNGReader()
img.SetFileName(os.path.normpath(image_path))
img.Update()
_extent = img.GetDataExtent()
ConstPixelDims = [
_extent[1] - _extent[0] + 1, _extent[3] - _extent[2] + 1,
_extent[5] - _extent[4] + 1
]
img_data = img.GetOutput()
datapointer = img_data.GetPointData()
assert (datapointer.GetNumberOfArrays() == 1)
vtkarray = datapointer.GetArray(0)
img = vtk.util.numpy_support.vtk_to_numpy(vtkarray)
img = img.reshape(ConstPixelDims, order='F')
img = img / np.max(img)
img = img.astype('float32')
#self.polar(img)
return img
def __len__(self):
"""Denotes the number of batches per epoch
:return: number of batches per epoch
"""
#return int(np.floor(len(self.list_IDs) / self.batch_size) * self.number_of_patches)
return int(np.floor(len(self.list_IDs)))
def __getitem__(self, index):
"""Generate one batch of data
:param index: index of the batch
:return: X and y when fitting. X only when predicting
"""
# Generate indexes of the batch
indexes = self.indexes[index * self.batch_size:(index + 1) *
self.batch_size]
# Find list of IDs
list_IDs_temp = [self.list_IDs[k] for k in indexes]
# Generate data
X = self._generate_X(list_IDs_temp)
if self.to_fit:
y = self._generate_y(list_IDs_temp)
return X, y
else:
return X
def _generate_y(self, list_IDs_temp):
"""Generates data containing batch_size images
:param list_IDs_temp: list of label ids to load
:return: batch of images
"""
Y = np.zeros(
(self.batch_size, *self.dimy, self.patch_size, self.n_channels))
# Generate data
for patch in list_IDs_temp:
for i, ID in enumerate(patch[1]):
path = self.mask_path + '/' + patch[0] + '/' + 'label_' + ID[
6:15] + '.png'
img = self._load_grayscale_image_VTK(path)[:, :, 0]
Y[0, :, :, i, 0] = img
if self.bool:
Y[0, :, :, i, :] = self.trans.apply_transform(
Y[0, :, :, i, :], self.param)
return Y
def _generate_X(self, list_IDs_temp):
"""Generates data containing batch_size images
:param list_IDs_temp: list of label ids to load
:return: batch of images
"""
X = np.zeros(
(self.batch_size, *self.dim, self.patch_size, self.n_channels))
self.param = self.trans.get_random_transform(self.dim)
if random.uniform(0, 1) >= 0.5:
self.bool = True
else:
self.bool = False
# Generate data
for patch in list_IDs_temp:
for i, ID in enumerate(patch[1]):
path = self.image_path + '/' + patch[0] + '/' + ID
img = self._load_grayscale_image_VTK(path)[:, :, 0]
X[0, :, :, i, 0] = img
if self.bool:
X[0, :, :, i, :] = self.trans.apply_transform(
X[0, :, :, i, :], self.param)
return X
def __next__(self):
if self.n >= self.max:
self.n = 0
result = self.__getitem__(self.n)
self.n += 1
return result
def on_epoch_end(self):
"""Updates indexes after each epoch
"""
self.indexes = np.arange(len(self.list_IDs))
if self.shuffle == True:
np.random.shuffle(self.indexes)
class WhiskerGenerator(keras.utils.Sequence):
def __init__(self,
baseDir,
timesteps,
batch_size,
rotation_range=0,
width_shift_range=0,
height_shift_range=0,
zoom_range=0,
mean=0.257,
std=0.288):
self.baseDir = baseDir
if not os.path.exists(baseDir):
raise ValueError("Invalid data path")
self.timesteps = timesteps
self.batch_size = batch_size
self.mean = np.reshape([mean], [1, 1, 1])
self.std = np.reshape([std], [1, 1, 1])
self.augment = rotation_range or width_shift_range or height_shift_range or zoom_range
if self.augment:
self.datagen = ImageDataGenerator(
rotation_range=rotation_range,
width_shift_range=width_shift_range,
height_shift_range=height_shift_range,
zoom_range=zoom_range)
extensions = {'png', 'PNG'}
self.sessions = []
for subdir in sorted(os.listdir(self.baseDir)):
if os.path.isdir(os.path.join(self.baseDir, subdir)):
self.sessions.append(subdir)
self.frameInfo = {}
for session in self.sessions:
if os.path.exists(
os.path.join(self.baseDir, session, "frameInfo.csv")):
self.frameInfo[session] = np.loadtxt(os.path.join(
self.baseDir, session, "frameInfo.csv"),
delimiter=",")
else:
print(session, "is not a valid session")
continue
self.training_examples = []
self.width = {}
for session in self.frameInfo:
counter = 0
for trial in self.frameInfo[session]:
if int(trial[1]) - int(trial[0]) < self.timesteps:
print("Trial not long enough, skipping...")
print("Frame:", trial[0])
continue
for frame in range(int(trial[0]),
int(trial[1]) - self.timesteps + 1):
self.onehot = np.zeros(self.timesteps + 1)
if int(trial[2]) - frame < self.timesteps and int(
trial[2]) - frame >= 0:
self.onehot[int(trial[2]) - frame] = 1
else:
self.onehot[self.timesteps] = 1
self.training_examples.append(
(session, frame, self.onehot))
counter += 1
self.width[session] = len(
os.listdir(os.path.join(self.baseDir,
session))[0].split(".")[0])
print("Found", counter, "images in", session)
self.on_epoch_end()
def __len__(self):
return int(np.ceil(len(self.training_examples) / self.batch_size))
def __getitem__(self, idx):
if idx >= int(np.ceil(len(self.training_examples) / self.batch_size)):
print("index too large for number of training examples")
return -1
self.x = [
list(
map(getImg, [
os.path.join(
self.baseDir, image[0],
str(frame).zfill(self.width[image[0]]) + ".png")
for frame in range(image[1], image[1] + self.timesteps)
]))
for image in self.training_examples[idx * self.batch_size:min(
(idx + 1) * self.batch_size, len(self.training_examples))]
]
if self.augment:
xfrms = [
self.datagen.get_random_transform(self.x[0][0].shape[:2])
for i in self.x
]
self.x = [[self.datagen.apply_transform(n, xfrms[i]) for n in x]
for i, x in enumerate(self.x)]
self.x = np.asarray(self.x)
self.x -= self.mean
self.x /= self.std
self.y = np.asarray([
ex[2] for ex in self.training_examples[idx * self.batch_size:min(
(idx + 1) * self.batch_size, len(self.training_examples))]
])
return self.x, self.y
def getClassWeights(self):
self.counts = np.zeros(self.timesteps + 1)
for example in self.training_examples:
self.counts[np.argmax(example[2])] += 1
self.counts = 1 / self.counts
self.scaling = (self.timesteps + 1) / np.sum(self.counts)
self.counts *= self.scaling
self.toReturn = {}
for i, weight in enumerate(self.counts):
self.toReturn[i] = weight
return self.toReturn
def on_epoch_end(self):
random.shuffle(self.training_examples)
class HeadPoseDataGenerator(Sequence):
'''
This class implements a basic Keras data generator overriding methods from its parent class Sequence. The purpose
of this data generator is to deliver in each batch a set of images from the subset used to initalize this generator
containing an equal number of members per class.
'''
def __init__(self, pose_dataframe, img_array, batch_size,
normalize=False, input_norm=None, tilt_norm=None, pan_norm=None,
augment=False, shift_range=None, zoom_range=None, brightness_range=None,
img_rescale=1, out_rescale=1):
'''
Initializes the data generator with the data from a given subset from the original dataset, and the values to
use when doing data augmentation.
Arguments:
pose_dataframe: Dataframe containing a list of each picture in the given subset and its pose values.
img_array: Numpy array containing the images from the given subset.
batch_size: Number of pictures per batch.
normalize: If the data shall be normalized or not.
input_norm: Tuple containing mean and std values for normalizing pictures in the dataset.
tilt_norm: Tuple containing mean and std values for normalizing tilt values in the dataset.
pan_norm: Tuple containing mean and std values for normalizing pan values in the dataset.
augment: If data augmentation shall be applied or not.
shift_range: Value (between 0 and 1) indicating the portion of the length of the side of each picture that
can be used to shift the picture (in both axes).
zoom_range: Tuple containing the minimum and maximum values used to apply zoom to each picture.
brightness_range: Tuple containing the minimum and maximum values used to apply a brightness transformation
to each picture.
img_rescale: Each pixel from every picture in the subset will be multiplied by this value.
out_rescale: Tilt and pan values for every picture in the subset will be multiplied by this value.
'''
# Create empty arrays for pictures and labels from the subset.
self.pics = []
self.labels = []
# Initialize batch size.
self.batch_size = batch_size
# Initialize normalization parameters.
self.normalize = normalize
self.input_norm = input_norm
'''
Initialize the parameter controlling if data augmentation shall be applied or not, and data augmentation
parameters.
'''
self.augment = augment
if self.augment == True:
self.generator = ImageDataGenerator(width_shift_range=shift_range, height_shift_range=shift_range,
brightness_range=brightness_range, zoom_range=zoom_range)
# Initialize scaling parameters.
self.img_rescale = img_rescale
self.out_rescale = out_rescale
'''
Initialize the iterator used to control the position of the next picture from every class that will be included
in a batch.
'''
self.common_iterator = 0
# Sort dataframe by class.
df = pose_dataframe.sort_values('class')
# Initialize the number of pictures in the dataset.
self.total_size = len(df.index)
# Load images and pose values into the previously created arrays.
prev_class = -1
class_index = -1
# For each image in the (ordered) dataset:
for index, row in df.iterrows():
'''
If the class for the current picture is different from the last class recorded, append an empty list for the
new class.
'''
if row['class'] != prev_class:
prev_class = row['class']
self.pics.append([])
self.labels.append([])
class_index = class_index + 1
# Append picture to corresponding class array.
self.pics[class_index].append(np.squeeze(img_array[index]))
# Append labels to corresponding class array (normalized and rescaled).
self.labels[class_index].append([((row['tilt'] * out_rescale) - tilt_norm[0]) / tilt_norm[1] , ((row['pan'] * out_rescale) - pan_norm[0]) / pan_norm[1]])
# Assert batch size is a multiple of the number of classes.
assert(batch_size % len(self.pics) == 0)
def __data_generation(self):
'''
Outputs a batch of pictures.
Returns:
X: Pictures in the batch.
y: Labels for each picture in the batch.
'''
# Create empty lists for pictures and labels.
X = []
y = []
# For each picture-per-class:
for i in range(int(self.batch_size / len(self.pics))):
# For each class:
for j in range(len(self.pics)):
# Select the next picture in the class list (start from beginning after the last picture).
pic = self.pics[j][int(self.common_iterator % len(self.pics[j]))]
pic = np.expand_dims(pic, axis=2)
# Apply data augmentation.
if self.augment == True:
transformation = self.generator.get_random_transform(pic.shape)
transformation['zx'] = transformation['zy']
pic = self.generator.apply_transform(pic, transformation)
# Rescale each pixel value in image.
pic = pic * self.img_rescale
# Normalize image.
if self.normalize == True:
pic = (pic - self.input_norm[0]) / self.input_norm[1]
# Add image and labels to the batch.
X.append(pic)
y.append(self.labels[j][int(self.common_iterator % len(self.labels[j]))])
# Update iterator.
self.common_iterator = self.common_iterator + 1
# Transform lists into Numpy arrays.
X = np.array(X)
y = np.array(y)
# Return images and labels.
return X, y
def __len__(self):
'''
Outputs the length (number of batches) that the data generator can provide.
Returns:
l: The length of the data generator.
'''
'''
Calculate the length of the data generator as the relation between the total number of images and the size of
each batch; in order to function properly with uneven class lengths the number is rounded to the smaller integer
bigger or equal to the obtained result.
'''
l = ceil(self.total_size / self.batch_size)
return l
def __getitem__(self, index):
'''
Outputs a new batch given the batch index.
Returns:
X: Pictures in the batch.
y: Labels for each picture in the batch.
'''
# Set the class iterator in the correct position for obtaining the requested batch.
self.common_iterator = index * int(self.batch_size / len(self.pics))
# Generate the batch.
X, y = self.__data_generation()
# Return images and labels for the requested batch.
return X, y
def reset(self):
'''
Resets the class iterator to its initial position.
'''
self.common_iterator = 0
class DataGenerator(tf.keras.utils.Sequence):
def __init__(
self,
directory,
batch_size=32,
target_size=(480, 640),
scale_size=1,
shuffle=True,
rotation_range=180,
zoom_range=[0.7, 1.5],
horizontal_flip=True,
vertical_flip=True,
fill_mode="nearest",
):
self.scale_size = scale_size
self.target_size = target_size
self.batch_size = batch_size
self.directory = directory
self.img_paths = []
self.img_paths_wo_ext = []
self.target = []
self.generator = ImageDataGenerator(
rotation_range=rotation_range,
zoom_range=zoom_range,
horizontal_flip=horizontal_flip,
vertical_flip=vertical_flip,
fill_mode=fill_mode,
)
for root, dirs, files in walk(directory):
for file in files:
if file.lower().endswith(".jpg") or file.lower().endswith(".png"):
self.img_paths.append(path.join(root, file))
self.img_paths_wo_ext.append(
path.splitext(path.join(root, file))[0]
)
elif file.lower().endswith(".txt"):
y = get_face_center_coordinates(root, "", file)
self.target.append(y)
self.targets = pd.DataFrame(self.target, columns=["x", "y"])
self.targets = self.targets.set_index(pd.Index(self.img_paths_wo_ext))
self.shuffle = shuffle
self.on_epoch_end()
def __len__(self):
return int(np.floor(len(self.img_paths) / self.batch_size))
def __getitem__(self, index):
indexes = self.indexes[index * self.batch_size : (index + 1) * self.batch_size]
list_paths = [self.img_paths[k] for k in indexes]
list_paths_wo_ext = [self.img_paths_wo_ext[k] for k in indexes]
X, y = self.__data_generation(list_paths, list_paths_wo_ext)
return X, y
def on_epoch_end(self):
self.indexes = np.arange(len(self.img_paths))
if self.shuffle == True:
np.random.shuffle(self.indexes)
def __data_generation(self, list_paths, list_paths_wo_ext):
width, height = self.target_size
scaled_width, scaled_height = (
int(width * self.scale_size),
int(height * self.scale_size),
)
X = np.empty(
(self.batch_size, scaled_width, scaled_height, 3), dtype=np.float32
)
y = [
[int(round(coordinate)) for coordinate in coordinates]
for coordinates in self.targets.loc[list_paths_wo_ext].values
]
for i, ID in enumerate(list_paths):
image = resize(imread(ID), (height, width))
empty_image_with_label = np.empty((width, height, 3), dtype=np.float32)
random_non_zero_point = [255, 36, 0]
empty_image_with_label[y[i][1]][y[i][0]] = np.array(random_non_zero_point)
while True:
transform_parameters = self.generator.get_random_transform(
image, seed=None
)
empty_image_with_label = self.generator.apply_transform(
empty_image_with_label, transform_parameters
)
transformed_label = np.nonzero(empty_image_with_label)
try:
if len(transformed_label[0]) != 0 and len(transformed_label[1]) != 0:
y[i][1] = transformed_label[0].mean() * self.scale_size
y[i][0] = transformed_label[1].mean() * self.scale_size
if isinstance(y[i][0], (int, float)) and isinstance(y[i][1], (int, float)) and y[i][0]!=float('nan') and y[i][1]!=float('nan'):
if not isinstance(y[i][0], bool) and not isinstance(y[i][1], bool) and y[i][0]*0 == 0 and y[i][1]*0 == 0:
break
except ValueError:
pass
X[i,] = self.generator.apply_transform(
resize(image, (scaled_height, scaled_width)), transform_parameters
)
return X, np.asarray(y)
unique.append(i)
while attr_count_crop[key] < crop:
ii = choice(unique)
arr = np.array(
skimage.io.imread(
img_path_bbox_attr_cls_tuples_list[ii][0].replace(
'\\', '/')))
w, h = arr.shape[1], arr.shape[0]
bbox = [
float(x) for x in img_path_bbox_attr_cls_tuples_list[ii]
[1].split('-')
]
if len(arr.shape) != 3 or arr.shape[2] != 3:
continue
transform_parameters = img_gen.get_random_transform(
arr.shape, seed=attr_count_crop[key])
res = img_gen.random_transform(arr, seed=attr_count_crop[key])
x = np.zeros((h, w, 3))
x[int(bbox[1] * h):int(bbox[3] * h) - 1,
int(bbox[0] * w):int(bbox[2] * w) - 1, 0] = 100
x = apply_affine_transform(
x,
transform_parameters.get('theta', 0),
transform_parameters.get('tx', 0),
transform_parameters.get('ty', 0),
transform_parameters.get('shear', 0),
transform_parameters.get('zx', 1),
transform_parameters.get('zy', 1),
row_axis=0,
col_axis=1,
channel_axis=2,
class NiftiImageIterator(Sequence):
"""
Niftiのパスから各軸から1枚ずつ画像を切り出して合計3枚のデータを作成する。
roiのパスも参照して乗算をすることで任意の位置だけ活用する。
正規化して学習データとして渡す。
kerasのSequenceクラスを継承している。
generatorで学習時に__getitem__()が叩かれて
バッチ分の学習データを生成して返すようになっている。
Args:
x_nifti_path (np.ndarray): nifti画像のパス。 shapeは(n,)
x_roi_path (np.ndarray): roi画像のパス。 shapeは(n,)
y (np.ndarray): ラベル。 shapeは(n, 2)
target_size (Tuple[int, int]): リサイズするときのサイズ。 (w, h)
ex_size (Tuple[int, int]): 画像をはっつけるキャンバスのサイズ。 (h, w)
test (bool): test用のgeneratorにするどうか。
preprocess_input (Callable): 前処理用の関数。汚いけどハードコーディングで渡してしまっている。
"""
def __init__(self,
x_nifti_path,
x_roi_path,
y,
target_size=(224, 224),
ex_size=(600, 600),
batch_size=32,
shuffle=False,
test=False,
preprocess_input=preprocess_input):
self.x_nifti_path = x_nifti_path
self.x_roi_path = x_roi_path
self.y = y
self.target_size = target_size
self.ex_size = ex_size
self.batch_size = batch_size
self.shuffle = shuffle
self.test = test
self.sample_num = len(self.y)
self.preprocess_input = preprocess_input
self.__get_exploration_order()
self.__create_data_gen()
def __getitem__(self, idx):
batch_ids = self.indexes[idx * self.batch_size:(idx + 1) *
self.batch_size]
x1, x2, x3, y = self.__data_generation(batch_ids)
return [x1, x2, x3], y
def __len__(self):
return int(np.ceil(len(self.y) / self.batch_size))
def __get_exploration_order(self):
self.indexes = np.arange(self.sample_num)
if self.shuffle:
self.indexes = np.random.shuffle(self.indexes)
def __create_data_gen(self):
if self.test:
self.datagen = None
else:
self.datagen = ImageDataGenerator(rotation_range=30,
horizontal_flip=True,
vertical_flip=True)
def __data_generation(self, batch_ids):
x1, x2, x3 = self.__get_imgs(batch_ids)
y = self.y[batch_ids]
return x1, x2, x3, y
def __get_imgs(self, batch_ids):
"""
画像を読み込んで諸々の処理をしてリスト形式で返す。
"""
x_nifti_path = self.x_nifti_path[batch_ids]
x_roi_path = self.x_roi_path[batch_ids]
img_x_list, img_y_list, img_z_list = [], [], []
for nii_path, roi_path in zip(x_nifti_path, x_roi_path):
nii = nib.load(nii_path)
box = nii.get_data()
nir = nib.load(roi_path)
roi = nir.get_data()
box = self.__normalize(box)
img_x, img_y, img_z = self.__get_slice(box, roi)
for img, img_list in zip([img_x, img_y, img_z],
[img_x_list, img_y_list, img_z_list]):
img = self.__resize_array(img)
img = self.preprocess_input(img)
if not self.test:
params = self.datagen.get_random_transform(img.shape)
img = self.datagen.apply_transform(img, params)
img /= 255.
img_list.append(img)
return np.asarray(img_x_list), np.asarray(img_y_list), np.asarray(
img_z_list)
def __normalize(self, arr):
"""
値域をいい感じにする。
値がハードコーディングで非常にきたない。
"""
arr[np.where(arr > 1024)] = 1024
arr[np.where(arr < -1024)] = -1024
arr = 255 * ((arr - arr.min()) / (arr.max() - arr.min()))
return arr
def __get_slice(self, arr, roi):
"""
各軸で切り出した際に断面積の最も大きい画像を返す。
"""
x_slice = np.zeros(self.ex_size)
y_slice = np.zeros(self.ex_size)
z_slice = np.zeros(self.ex_size)
best_x = np.argmax(np.sum(np.sum(roi, axis=1), axis=1))
best_y = np.argmax(np.sum(np.sum(roi, axis=0), axis=1))
best_z = np.argmax(np.sum(np.sum(roi, axis=0), axis=0))
a_s = arr.shape
x_slice[:a_s[1], :a_s[2]] = arr[best_x, :, :] * roi[best_x, :, :]
y_slice[:a_s[0], :a_s[2]] = arr[:, best_y, :] * roi[:, best_y, :]
z_slice[:a_s[0], :a_s[1]] = arr[:, :, best_z] * roi[:, :, best_z]
x_slice = np.stack([x_slice for _ in range(3)], axis=2)
y_slice = np.stack([y_slice for _ in range(3)], axis=2)
z_slice = np.stack([z_slice for _ in range(3)], axis=2)
return x_slice, y_slice, z_slice
def __resize_array(self, img):
pilimg = Image.fromarray(np.uint8(img))
pilimg = pilimg.resize(self.target_size)
return np.asarray(pilimg).astype(np.float32)
class DataGenerator2(DataGenerator):
"""Generates data for Keras
Sequence based data generator. Suitable for building data generator for training and prediction.
"""
def __init__(self,
list_IDs,
image_path,
mask_path,
to_fit=True,
batch_size=32,
dim=(512, 512),
n_channels=1,
n_classes=10,
shuffle=True,
data_gen_args=None):
"""Initialization
:param list_IDs: list of all 'label' ids to use in the generator
:param image_path: path to images location
:param mask_path: path to masks location
:param to_fit: True to return X and y, False to return X only
:param batch_size: batch size at each iteration
:param dim: tuple indicating image dimension
:param n_channels: number of image channels
:param n_classes: number of output masks
:param shuffle: True to shuffle label indexes after every epoch
"""
self.bool = False
if data_gen_args != None:
self.trans = ImageDataGenerator(**data_gen_args)
super().__init__(list_IDs, image_path, mask_path, to_fit, batch_size,
dim, n_channels, n_classes, shuffle)
def _generate_X(self, list_IDs_temp):
"""Generates data containing batch_size images
:param list_IDs_temp: list of label ids to load
:return: batch of images
"""
# Initialization
X = np.empty((self.batch_size, *self.dim, self.n_channels))
self.param = self.trans.get_random_transform(self.dim)
if random.uniform(0, 1) >= 0.5:
self.bool = True
else:
self.bool = False
# Generate data
for i, ID in enumerate(list_IDs_temp):
# Store sample
X[i, ] = self._load_dicom_image(self.image_path + '/' + ID)
# X[i,] = self.apply_transform(X[i,],self.get_random_transform((1,512,512)))
if self.bool:
X[i, ] = self.trans.apply_transform(X[i, ], self.param)
# X=np.expand_dims(X, 4)
return X
def _generate_y(self, list_IDs_temp):
"""Generates data containing batch_size masks
:param list_IDs_temp: list of label ids to load
:return: batch if masks
"""
y = np.empty((self.batch_size, *self.dim, self.n_channels))
# Generate data
for i, ID in enumerate(list_IDs_temp):
# Store sample
y[i, ] = self._load_grayscale_image_VTK(self.mask_path + '/' +
'label_' + ID[6:15] +
'.png')
if self.bool:
y[i, ] = self.trans.apply_transform(y[i, ], self.param)
return y
