图像修复Image Inpainting数据集制作

目录

掩码数据集

原始图像数据集

制作过程

1 将所有图片统一裁剪至相同尺寸（512×512）

2 顺序命名

3 将mask覆盖至原图上以达成破损图像制作

结果展示

由于本人目前所研究的大课题是GAN图像修复，但小方向做的具体修复项目目前并没有合适的数据集，因此只能自己制作，现把制作步骤写出，方便大家一起交流学习~

掩码数据集

所用掩码数据集选自于 Liu 等人2018年发布的论文 Image Inpainting for Irregular Holes Using Partial Convolutions。该论文中提出用部分卷积解决inpainting的任务的同时，也公布了一个大型的mask数据集。

下载链接：Image Inpainting for Irregular Holes Using Partial Convolutions - NVIDIA ADLR

打开后如下图所示

本人选用的是Testing Set中所使用的掩码数据集，共包含12000张mask图像，一般来说自己制作数据集的话这个数量也是够用的。

数据集内包含大、中、小三种覆盖面积的mask，可根据具体需求筛选最合适自己项目的。

原始图像数据集

根据具体需求制作，一般数量在1~2k张即可。

制作过程

1 将所有图片统一裁剪至相同尺寸（512×512）

由于上述的mask数据集掩码大小为512×512，为方便后续操作，我们在进行破损图像制作时，先统一将所有图片裁剪至相同尺寸。代码如下：

from PIL import Image
import os

def resize_images(input_folder, output_folder, size=(512, 512)):
    # 创建输出文件夹如果它不存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 初始化计数器
    counter = 1

    # 遍历输入文件夹中的所有文件
    for filename in os.listdir(input_folder):
        if filename.endswith(".jpg"):
            # 构建完整的文件路径
            input_path = os.path.join(input_folder, filename)
            output_path = os.path.join(output_folder, f"{counter}.jpg")  # 新的文件名

            # 打开图片并调整大小
            with Image.open(input_path) as img:
                resized_img = img.resize(size, Image.Resampling.LANCZOS)
                resized_img.save(output_path)

            # 更新计数器
            counter += 1

# 设置输入和输出文件夹路径
input_folder = '输入的初始图片文件夹路径'
output_folder = '输出的512×512大小图片文件夹路径'

# 调用函数执行转换
resize_images(input_folder, output_folder)

2 顺序命名

由于不可保证数据集内的图片命名一致，为了规范操作，在经过裁剪后，我们再对所有图片进行顺序重命名操作（mask与完整图像文件夹都需要）。代码如下：

import os

def rename_images(folder_path):
    # 创建一个计数器
    counter = 1

    # 遍历文件夹中的所有文件
    for filename in os.listdir(folder_path):
        if filename.endswith(".png"):
            # 构建完整的文件路径
            old_file_path = os.path.join(folder_path, filename)
            # 新的文件名
            new_filename = f"{counter}.jpg"
            # 构建新的文件路径
            new_file_path = os.path.join(folder_path, new_filename)

            # 重命名文件
            os.rename(old_file_path, new_file_path)

            # 更新计数器
            counter += 1

# 设置文件夹路径
folder_path = '需要顺序命名图片的文件夹路径'

# 调用函数执行重命名
rename_images(folder_path)

3 将mask覆盖至原图上以达成破损图像制作

在制作好所需的mask文件夹以及对应的初始图像文件夹后（注意二者内图像数量与名称一一对应），下述代码可根据掩码中的白色区域（即破损区域）来修改原始图像中的对应像素，从而模拟破损的效果，并将处理后的图像保存到指定的输出文件夹中。

from PIL import Image
import numpy as np
import os

def apply_damage_effect(image_folder, mask_folder, output_folder):
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历图像文件夹中的所有文件
    for image_file in os.listdir(image_folder):
        if image_file.endswith('.jpg'):
            image_path = os.path.join(image_folder, image_file)
            mask_path = os.path.join(mask_folder, image_file)  # 假设掩码文件与图像文件同名

            # 检查掩码文件是否存在
            if not os.path.exists(mask_path):
                print(f"Mask file not found for {image_file}, skipping...")
                continue

            try:
                # 加载图像和掩码
                image = Image.open(image_path).convert('RGB')
                mask = Image.open(mask_path).convert('L')  # 转换为灰度图像

                # 将图像和掩码转换为 NumPy 数组
                image_data = np.array(image)
                mask_data = np.array(mask)

                # 创建模拟破损的效果
                # 假设白色部分（高值）是破损，我们将其设置为白色
                damage_color = [255, 255, 255]  # RGB白色
                mask_indices = mask_data > 128  # 假定大于128的像素值为白色部分
                image_data[mask_indices] = damage_color  # 应用破损效果

                # 将修改后的数据转换回图像
                damaged_image = Image.fromarray(image_data)

                # 保存结果图像
                output_path = os.path.join(output_folder, image_file)
                damaged_image.save(output_path, 'JPEG')
                print(f"Processed and saved {output_path}")
            except Exception as e:
                print(f"Failed to process {image_file}: {e}")

# 设置文件夹路径
image_folder = '统一尺寸与名称原图的文件夹路径'
mask_folder = '选出的mask文件夹'
output_folder = '制作后的破损图像文件夹'

# 调用函数
apply_damage_effect(image_folder, mask_folder, output_folder)

结果展示

通过上述步骤，我们就可以得到一般图像修复任务中需要的完整图像以及破损图像数据集了。

masked图像

original图像