最近在玩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 多平台发布