共计 4361 个字符，预计需要花费 11 分钟才能阅读完成。
最近在玩 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 多平台公布
正文完
