自定义数据增强生成图片数量未达预期问题排查

问题描述 #

期望通过自定义数据增强代码生成5700张图片（原始图片数量的10倍），但实际仅生成了1140张（仅为原始数量的2倍），需排查代码问题。

原始代码 #

# Custom augmentation functions
def random_translate(image, max_translate):
    shift = np.random.uniform(-max_translate, max_translate, 3)
    translated_image = scipy.ndimage.shift(image, shift, mode='nearest')
    return translated_image

def random_rotate(image, max_angle):
    angles = np.random.uniform(-max_angle, max_angle, 3)
    rotated_image = scipy.ndimage.rotate(image, angles[0], axes=(1, 2), reshape=False)
    rotated_image = scipy.ndimage.rotate(rotated_image, angles[1], axes=(0, 2), reshape=False)
    rotated_image = scipy.ndimage.rotate(rotated_image, angles[2], axes=(0, 1), reshape=False)
    return rotated_image

def random_flip(image):
    if random.random() > 0.5:
        image = np.flip(image, axis=0)
    if random.random() > 0.5:
        image = np.flip(image, axis=1)
    if random.random() > 0.5:
        image = np.flip(image, axis=2)
    return image

def random_noise(image, noise_level=0.01):
    noise = np.random.normal(0, noise_level, image.shape)
    noisy_image = image + noise
    return noisy_image

def random_brightness(image, max_delta=0.2):
    delta = np.random.uniform(-max_delta, max_delta)
    bright_image = np.clip(image + delta, 0, 1)  # Clip to maintain valid pixel range
    return bright_image

def random_contrast(image, lower=0.8, upper=1.2):
    factor = np.random.uniform(lower, upper)
    mean = np.mean(image, axis=(0, 1, 2), keepdims=True)
    contrast_image = np.clip((image - mean) * factor + mean, 0, 1)
    return contrast_image

def random_scale(image, min_scale=0.9, max_scale=1.1):
    scale = np.random.uniform(min_scale, max_scale)
    height, width = image.shape[:2]
    scaled_image = scipy.ndimage.zoom(image, (scale, scale, 1), order=1)
    if scale < 1.0:
        pad_height = (height - scaled_image.shape[0]) // 2
        pad_width = (width - scaled_image.shape[1]) // 2
        scaled_image = np.pad(scaled_image, ((pad_height, pad_height), (pad_width, pad_width), (0, 0)), mode='constant')
    else:
        start_height = (scaled_image.shape[0] - height) // 2
        start_width = (scaled_image.shape[1] - width) // 2
        scaled_image = scaled_image[start_height:start_height + height, start_width:start_width + width]
    return scaled_image

def random_shear(image, max_shear=0.2):
    shear = np.random.uniform(-max_shear, max_shear)
    afine_tf = tf.keras.preprocessing.image.random_shear(shear)
    shear_image = tf.keras.preprocessing.image.apply_affine_transform(image, shear=afine_tf)
    return shear_image

def elastic_transform(image, alpha=1000, sigma=30):
    random_state = np.random.RandomState(None)
    shape = image.shape
    dx = scipy.ndimage.gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
    dy = scipy.ndimage.gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
    dz = scipy.ndimage.gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
    x, y, z = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), np.arange(shape[2]), indexing='ij')
    indices = np.reshape(x + dx, (-1, 1)), np.reshape(y + dy, (-1, 1)), np.reshape(z + dz, (-1, 1))
    distored_image = scipy.ndimage.map_coordinates(image, indices, order=1, mode='reflect')
    return distored_image.reshape(image.shape)


def custom_data_generator(images, labels, batch_size,
                          max_translate=10, max_angle=15, noise_level=0.01,
                          max_delta=0.2, lower_contrast=0.8, upper_contrast=1.2,
                          min_scale=0.9, max_scale=1.1, max_shear=0.2, alpha=1000, sigma=30):
    num_images = images.shape[0]
    while True:
        batch_indices = np.random.choice(num_images, batch_size)
        batch_images = []
        batch_labels = []
        for idx in batch_indices:
            image = images[idx]
            label = labels[idx]

            # Apply all augmentation functions
            image = random_translate(image, max_translate)
            image = random_rotate(image, max_angle)
            image = random_flip(image)
            image = random_noise(image, noise_level)
            image = random_brightness(image, max_delta)
            image = random_contrast(image, lower_contrast, upper_contrast)
            image = random_scale(image, min_scale, max_scale)
            image = tf.keras.preprocessing.image.random_shear(image, max_shear)
            image = elastic_transform(image, alpha, sigma)

            batch_images.append(image)
            batch_labels.append(label)

        batch_images = np.array(batch_images)
        batch_labels = np.array(batch_labels)
        yield batch_images, batch_labels


# Directory to save the images
save_dir = 'mri_augmented_images'
os.makedirs(save_dir, exist_ok=True)

for i in range(mri_resized.shape[0]):
    for j in range(mri_resized.shape[3]):  # Loop over channels
        img_array = mri_resized[i, :, :, j]
        img = Image.fromarray((img_array * 255).astype('uint8'))  # Scale to [0, 255] and convert to uint8
        img.save(os.path.join(save_dir, f'image_{i}_channel_{j}.png'))

# Using the custom data generator to augment images and retain labels
target_num_images = 5700
batch_size = 3  # Define your batch size
num_batches_needed = (target_num_images + batch_size - 1) // batch_size
data_gen = custom_data_generator(mri_resized, labels, batch_size=batch_size)

# Generate and save augmented images with labels
for batch_images, batch_labels in data_gen:  # Directly iterate over the generator
    for j, img_array in enumerate(batch_images):
        img = Image.fromarray((img_array[:, :, 0] * 255).astype('uint8'))  # Save only the first channel
        img_index = i * batch_size + j
        img.save(os.path.join(save_dir, f'aug_image_{img_index}_label_{batch_labels[j]}.png'))

print(f"Generated and saved {target_num_images} augmented images.")

问题排查与修正 #

1. 生成循环无终止条件+索引计算错误 #

原始代码直接遍历无限生成器data_gen，且使用保存原始图循环的残留变量i计算img_index，导致所有增强图索引重复，文件被不断覆盖，最终仅保留少量图片。

修正方案：跟踪已生成图片数量，达到目标后终止循环，用独立计数器作为索引：

# 修正后的增强图生成与保存代码
generated_count = 0
data_gen = custom_data_generator(mri_resized, labels, batch_size=batch_size)

while generated_count < target_num_images:
    batch_images, batch_labels = next(data_gen)
    # 取当前batch中需要的数量（避免最后一个batch超出目标）
    take = min(len(batch_images), target_num_images - generated_count)
    for j in range(take):
        img_array = batch_images[j]
        img = Image.fromarray((img_array[:, :, 0] * 255).astype('uint8'))
        img.save(os.path.join(save_dir, f'aug_image_{generated_count}_label_{batch_labels[j]}.png'))
        generated_count += 1

print(f"Generated and saved {generated_count} augmented images.")

2. random_shear函数用法错误 #

原函数中tf.keras.preprocessing.image.random_shear调用参数错误，该函数第一个参数应为图像而非shear值，会导致运行报错中断生成。

修正方案：

def random_shear(image, max_shear=0.2):
    shear = np.random.uniform(-max_shear, max_shear)
    # 直接调用random_shear处理图像，intensity参数传入剪切角度
    shear_image = tf.keras.preprocessing.image.random_shear(image, intensity=shear)
    return shear_image

3. 2D图像误用3D变换逻辑 #

random_translate、random_rotate、elastic_transform函数针对3D图像设计了三维变换（如shift取3个值、三维旋转、dz分量），但你的图像是2D+通道结构（shape为(height, width, channels)），通道维度不需要变换，会导致无效操作甚至图像变形。

修正示例（以random_translate为例）：

def random_translate(image, max_translate):
    # 仅对height和width维度做平移，通道维度不变
    shift = np.random.uniform(-max_translate, max_translate, 2)
    # 补全通道维度的平移量为0
    shift = np.append(shift, 0)
    translated_image = scipy.ndimage.shift(image, shift, mode='nearest')
    return translated_image

4. 原始图与增强图数量混淆 #

原始代码先保存了所有原始图像的通道图，这部分会被计入总文件数，但你需要的是5700张增强图，需确保增强图的生成计数独立于原始图。

总结 #

核心问题是增强图生成循环无终止逻辑且索引重复导致文件覆盖，同时存在变换函数的参数错误。修正上述问题后，即可生成目标数量的增强图片。

内容的提问来源于stack exchange，提问作者dhruv Gangwani