具有数据增强功能的 keras 中 3D 体积的图像生成器

Question

由于 keras 的 ImageDataGenerator 不适合 3D 体积，我开始为 keras 编写自己的生成器（语义分割，而不是分类！）。

1) 如果有人已经修改了 ImageDataGenerator 代码以处理 3D 体积，请分享它！这家伙已经为视频做到了。

2）根据本教程，我编写了一个自定义生成器。

import glob
import os

import keras
import numpy as np
import skimage
from imgaug import augmenters as iaa


class DataGenerator(keras.utils.Sequence):
    """Generates data for Keras"""
    """This structure guarantees that the network will only train once on each sample per epoch"""

    def __init__(self, list_IDs, im_path, label_path, batch_size=4, dim=(128, 128, 128),
                 n_classes=4, shuffle=True, augment=False):
        'Initialization'
        self.dim = dim
        self.batch_size = batch_size
        self.list_IDs = list_IDs
        self.im_path = im_path
        self.label_path = label_path
        self.n_classes = n_classes
        self.shuffle = shuffle
        self.augment = augment
        self.on_epoch_end()

    def __len__(self):
        'Denotes the number of batches per epoch'
        return int(np.floor(len(self.list_IDs) / self.batch_size))

    def __getitem__(self, index):
        'Generate one batch of data'
        # 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, y = self.__data_generation(list_IDs_temp)

        return X, y

    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)

    def __data_generation(self, list_IDs_temp):

        if self.augment:
            pass

        if not self.augment:

            X = np.empty([self.batch_size, *self.dim])
            Y = np.empty([self.batch_size, *self.dim, self.n_classes])

            # Generate data
            for i, ID in enumerate(list_IDs_temp):
                img_X = skimage.io.imread(os.path.join(im_path, ID))
                X[i,] = img_X

                img_Y = skimage.io.imread(os.path.join(label_path, ID))
                Y[i,] = keras.utils.to_categorical(img_Y, num_classes=self.n_classes)

            X = X.reshape(self.batch_size, *self.dim, 1)
            return X, Y

params = {'dim': (128, 128, 128),
          'batch_size': 4,
          'im_path': "some/path/for/the/images/",
          'label_path': "some/path/for/the/label_images",
          'n_classes': 4,
          'shuffle': True,
          'augment': True}



partition = {}
im_path = "some/path/for/the/images/"
label_path = "some/path/for/the/label_images/"

images = glob.glob(os.path.join(im_path, "*.tif"))
images_IDs = [name.split("/")[-1] for name in images]

partition['train'] = images_IDs

training_generator = DataGenerator(partition['train'], **params)

我的图像有大小(128, 128, 128)，当我加载它们时，我会得到一个 5D 张量(batch_size, depth, heigt, width, number_of_channels)，例如(4, 128, 128, 128, 1). 对于 label_images（它们具有相同的尺寸并且是单通道编码（值 1 = 标签 1，值 2 = 标签 2，值 3 = 标签 3 和值 0 = 标签 4 或背景））我得到了标签的二进制表示使用to_categorical()来自 keras的函数并以 5D 结束，例如(4, 128, 128, 128, 4). images 和 label_images 具有相同的名称并位于不同的文件夹中。

由于我只有很少的图像，我想通过图像增强来扩展图像的总数。我将如何用这个生成器做到这一点？我已经成功测试了这个imgaug包，但我没有将图像添加到我的集合中，而是只转换现有图像（例如水平翻转它们）

编辑：我对数据增强有误解。请参阅有关图像增强的这篇文章。图像将通过随机转换（即时）传递。现在我只需要收集足够的数据并使用imgaug. 我会尽快更新这个。

Answer 1

我找到了一个用于 3D 体积的 Keras customDataGenerator 实现。这是GitHub 链接。该实现可以轻松扩展以包含新的增强技术。这是我在我的项目（3D 体积语义分割）中工作的一个最小工作示例，基于我在链接中共享的实现：

from generator import customImageDataGenerator

def generator(images, groundtruth, batch):
    """Load a batch of augmented images"""
  
    gen = customImageDataGenerator(mirroring=True, 
                                   rotation_90=True,
                                   transpose_axes=True
                                   )

    for b in gen.flow(x=images, y=groundtruth, batch_size=batch):
        yield (b[0], (b[1]).astype(float))

# images = (123, 48,48,48,1)
# groundtruth = (123, 48,48,48,1)
history = model.fit(
    x=generator(images, groundtruth, batchSize),
    validation_data=(imagesTest, groundtruthTest),
    steps_per_epoch=len(images) / batchSize,
    epochs=epochs,
    callbacks=[callbacks],
    )