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