[图片训练]处理图片脚本

最近在玩AI绘图,进行Lora模型或者其他训练时需要对图片进行预处理,当然stable diffusion web ui本身就有预处理脚本了,lora训练库里也有相关脚本.

但我这里还是写了一个分享一下

import os
import argparse
from PIL import Image,ImageFilter
import imghdr
from rembg import remove


imgType_list = {'jpg','bmp','png','jpeg','rgb','tif',"webp"}
global outputpath
global autopath
def save(img,file):
    """
    save img
    :param img:
    :param file:
    :return:
    """
    if os.path.isfile(outputpath):
        if autopath:
            img.save(file)
        else:
            img.save(outputpath)
    else:
        if autopath:
            img.save(file)
        else:
            img.save(os.path.join(outputpath,os.path.basename(file)))

def resizeimg(files,size):
    """
    resize img to size
    :param file:
    :return:
    """
    print("图片进行Resize")
    if os.path.isfile(files):
        img = Image.open(files)
        img = img.resize(size)
        save(img,files)
    else:
        for root, dirs, file in os.walk(files):
            for f in file:
                if imghdr.what(os.path.join(root, f)) in imgType_list:
                    img = Image.open(os.path.join(root, f))
                    img = img.resize(size)
                    save(img, os.path.join(root, f))
            for dir in dirs:
                resizeimg(os.path.join(root,dir),size)
def renameimg(files):
    """
    rename img index
    :param file:
    :return:
    """
    print("图片进行Rename")
    if os.path.isfile(files):
        img = Image.open(files)
        if autopath:
            img.save(os.path.join(os.path.dirname(files),"0.jpg"))
        else:
            img.save(os.path.join(outputpath, "0.jpg"))
    else:
        for root,dirs, files in os.walk(files):
            for i,f in enumerate(files):
                if imghdr.what(os.path.join(root, f)) in imgType_list:
                    img = Image.open(os.path.join(root, f))
                    if autopath:
                        img.save(os.path.join(root,str(i)+".jpg"))
                    else:
                        img.save(os.path.join(outputpath,str(i)+".jpg"))
            for dir in dirs:
                renameimg(os.path.join(root,dir))

def removeimgbg(files):
    """
    remove img background
    :param file:
    :return:
    """
    print("图片进行去除背景")
    if os.path.isfile(files):
        img = Image.open(files)
        if files.split('.')[-1] == "jpg":
            img = img.convert('RGB')
            img = img.filter(ImageFilter.BLUR)
        output = remove(img)
        if autopath:
            output.save(os.path.join(os.path.dirname(files),os.path.basename(files).split('.')[0]+".png"))
        else:
            output.save(os.path.join(os.path.dirname(outputpath),os.path.basename(outputpath).split('.')[0]+".png"))
    else:
        for file in os.listdir(files):
            img = Image.open(os.path.join(files,file))
            output = remove(img)
            if autopath:
                output.save(os.path.join(files, os.path.basename(file).split('.')[0]+".png"))
            else:
                output.save(os.path.join(outputpath, os.path.basename(file).split('.')[0]+".png"))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', help='input file path,allow file and dir',required=True) # get input path
    parser.add_argument('-o', '--output', help='output file path') # get output path
    parser.add_argument('-s', '--size', help='output size',nargs=2,type=int)  # get output size
    parser.add_argument('-r', '--removebg', help='remove background',type=bool) # remove background
    parser.add_argument('-e', '--rename', help='rename file',type=bool) # rename file (index)
    parser.add_argument('-a', '--autopath', help='outputpath same as the img path',type=bool) # rename file (index)
    

    args = parser.parse_args()
    file = args.input
    autopath = args.autopath
    if os.path.isfile(file):
        outputpath = args.output if args.output else file
    elif os.path.isdir(file):
        outputpath = args.output if args.output else file
        if not os.path.exists(outputpath):
            os.makedirs(outputpath)
    else:
        print("input file or dir is not exist")
        exit()
        
    size = args.size
    removebg = args.removebg
    rename = args.rename


    if os.path.exists(file):
        """
        if input is a file or dir 
        """
        if size:
            resizeimg(file,size=size)
        if removebg:
            removeimgbg(file)
        if rename:
            renameimg(file)
    else:
        print("input file or dir is not exist")

介绍

主要是图片剪裁,重命名和去除背景.去除背景也是最重要的,使用了rembg这个包.

官网danielgatis/rembg: Rembg is a tool to remove images background (github.com),可以考虑GPU版本.

同时测试时也发现如果是.jpg文件去掉背景会存在一些问题,是什么通道还是啥问题,不过使用格式转换啥的最终也解决了.

包环境

requirements如下,因为环境问题,可能有一些用不上的包.

aiohttp==3.8.4
aiosignal==1.3.1
anyio==3.6.2
async-timeout==4.0.2
asyncer==0.0.2
attrs==23.1.0
certifi==2022.12.7
charset-normalizer==3.1.0
click==8.1.3
colorama==0.4.6
coloredlogs==15.0.1
fastapi==0.95.1
filelock==3.12.0
filetype==1.2.0
flatbuffers==23.3.3
frozenlist==1.3.3
h11==0.14.0
humanfriendly==10.0
idna==3.4
ImageHash==4.3.1
imageio==2.27.0
Jinja2==3.1.2
lazy_loader==0.2
llvmlite==0.39.1
MarkupSafe==2.1.2
mpmath==1.3.0
multidict==6.0.4
networkx==3.1
numba==0.56.4
numpy==1.23.5
onnxruntime==1.14.1
opencv-python-headless==4.7.0.72
packaging==23.1
Pillow==9.5.0
platformdirs==3.2.0
pooch==1.7.0
protobuf==4.22.3
pydantic==1.10.7
PyMatting==1.1.8
pyreadline3==3.4.1
python-multipart==0.0.6
rembg==2.0.32
requests==2.28.2
scikit-image==0.20.0
scipy==1.10.1
sniffio==1.3.0
starlette==0.26.1
sympy==1.11.1
tifffile==2023.4.12
torch==2.0.0
tqdm==4.65.0
typing_extensions==4.5.0
urllib3==1.26.15
uvicorn==0.21.1
watchdog==3.0.0
yarl==1.9.1

本文由mdnice多平台发布