Stable Diffusion训练图片时，简陋的数据处理

0 图片从命名

如果有强迫症，看到似乎乱码的命名会不舒服，那么就批量从命名

import os


def rename_files_in_directory(directory, key_word, new_suffix):
    i = 1
    for filename in os.listdir(directory):
        new_file = key_word + str(i).zfill(3) + new_suffix
        source = os.path.join(directory, filename)
        destination = os.path.join(directory, new_file)
        os.rename(source, destination)
        i += 1
# 使用方法
# rename_files_in_directory('/path/to/directory', '.new_suffix')
# D:\SdTrainerGUI\lora-scripts-v1.7.3\train\XiboBird\5_zkz

1 批量缩小图片分辨率

如果是用同一个相机拍的，分辨率都是一样的，只不过分辨率太大了8K以上的分辨率显卡受不了

from PIL import Image
import os


def resize_image(image_path, output_path, scale_factor):
    # 打开图片
    img = Image.open(image_path)

    # 获取图片的宽度和高度
    width, height = img.size

    # 计算新的宽度和高度
    new_width = width // scale_factor
    new_height = height // scale_factor

    # 使用ANTIALIAS滤镜来缩小图片
    # new_img = img.resize((new_width, new_height), Image.ANTIALIAS)
    new_img = img.resize((new_width, new_height), Image.ANTIALIAS)

    # 保存新图片
    new_img.save(output_path)


def get_all_image(path, file_extension=".jpg"):
    return [os.path.join(path, f) for f in os.listdir(path) if f.endswith(file_extension)]


def process_images(catalog_of_original_images, file_extension, scale_factor):
    # 获取原始图像目录的上一级目录
    parent_directory = os.path.dirname(catalog_of_original_images)

    # 设置输出目录
    output_catalog = os.path.join(parent_directory, "output")

    # 创建输出目录
    if not os.path.exists(output_catalog):
        os.makedirs(output_catalog)

    image_list = get_all_image(catalog_of_original_images, file_extension)

    for image in image_list:
        # 获取图片的文件名
        image_name = os.path.basename(image)
        # 设置输出图片的路径
        output_image_path = os.path.join(output_catalog, image_name)
        # 缩小图片并保存
        resize_image(image, output_image_path, scale_factor)


if __name__ == '__main__':
    process_images(r"E:\Dwk\Photos\祥春鸟", ".jpg", 10)

2 图片数据集增强

最简易的增强是图片镜像，就是左右颠倒各一张，图片数据集数量直接翻倍

import os
from PIL import Image, ImageOps


def data_enhancement(image_catalog, file_extension=".jpg"):
    image_list = [os.path.join(image_catalog, f) for f in os.listdir(image_catalog) if f.endswith(file_extension)]
    for image in image_list:
        # 打开图片
        img = Image.open(image)
        # 创建镜像图片
        mirror_img = ImageOps.mirror(img)
        # 获取图片的文件名（不包括后缀）
        image_name = os.path.splitext(os.path.basename(image))[0]
        # 设置镜像图片的文件名
        mirror_image_name = image_name + "_mirror" + file_extension
        # 设置镜像图片的路径
        mirror_image_path = os.path.join(image_catalog, mirror_image_name)
        # 保存镜像图片
        mirror_img.save(mirror_image_path)


if __name__ == '__main__':
    data_enhancement(r"E:\Dwk\Photos\output", ".jpg")

3 tag内容批量修改（这里是只替换）

避免一个个文件打开逐个tag修改

import os


def replace_words_in_files(directory, old_word, new_word):
    # 获取指定目录下的所有文件
    files = os.listdir(directory)

    # 遍历所有文件
    for file in files:
        # 检查文件是否为.txt文件
        if file.endswith('.txt'):
            # 构建完整的文件路径
            file_path = os.path.join(directory, file)

            # 打开文件
            with open(file_path, 'r') as f:
                content = f.read()

            # 替换内容
            content = content.replace(old_word, new_word)

            # 写回文件
            with open(file_path, 'w') as f:
                f.write(content)


if __name__ == '__main__':
    replace_words_in_files(r'D:\SdTrainerGUI\lora-scripts-v1.7.3\train\PreprocessingOutput','girl','boy')