最近在玩AI绘图,进行Lora模型或者其余训练时须要对图片进行预处理,当然stable diffusion web ui自身就有预处理脚本了,lora训练库里也有相干脚本.
但我这里还是写了一个分享一下
import osimport argparsefrom PIL import Image,ImageFilterimport imghdrfrom rembg import removeimgType_list = {'jpg','bmp','png','jpeg','rgb','tif',"webp"}global outputpathglobal autopathdef 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.4aiosignal==1.3.1anyio==3.6.2async-timeout==4.0.2asyncer==0.0.2attrs==23.1.0certifi==2022.12.7charset-normalizer==3.1.0click==8.1.3colorama==0.4.6coloredlogs==15.0.1fastapi==0.95.1filelock==3.12.0filetype==1.2.0flatbuffers==23.3.3frozenlist==1.3.3h11==0.14.0humanfriendly==10.0idna==3.4ImageHash==4.3.1imageio==2.27.0Jinja2==3.1.2lazy_loader==0.2llvmlite==0.39.1MarkupSafe==2.1.2mpmath==1.3.0multidict==6.0.4networkx==3.1numba==0.56.4numpy==1.23.5onnxruntime==1.14.1opencv-python-headless==4.7.0.72packaging==23.1Pillow==9.5.0platformdirs==3.2.0pooch==1.7.0protobuf==4.22.3pydantic==1.10.7PyMatting==1.1.8pyreadline3==3.4.1python-multipart==0.0.6rembg==2.0.32requests==2.28.2scikit-image==0.20.0scipy==1.10.1sniffio==1.3.0starlette==0.26.1sympy==1.11.1tifffile==2023.4.12torch==2.0.0tqdm==4.65.0typing_extensions==4.5.0urllib3==1.26.15uvicorn==0.21.1watchdog==3.0.0yarl==1.9.1本文由mdnice多平台公布