一直在关注U^2-Net 显著性检测算法,可以用来静态背景分割。实践中可以把照片生成素描。
由于该算法需要使用到GPU计算,在华为云上实践了下,非常奈斯,不过却花费了4大洋,于是准备在本机跑一下,go。
1 运行环境准备
1.1 显卡硬件准备
运行该算法我们需要建立一个合适的运行环境,首先我们需要一个Nvidia的显卡,可以通过nvidia-smi在cmd窗口中看到相关的信息。
1.2 安装cuda和cudnn的安装
我们需要安装CUDA和CUDNN,注意根据最新的驱动版本要求,cuda需要11.6以上。cudnn需要注册才能下载,按提示安装即可。
1.3 pytorch安装
一般可以在pytorch网站(PyTorch)根据使用环境生成命令行,需要选择GPU/CPU、cuda版本等。自己pip安装pytorch时一直提示pytorch的名字其实是torch,然后改了也不能运行,非常痛苦。
蓝色即安装命令行。
1.4 其他依赖包安装
该算法仅依赖很少的包:如 Pytorch、Numpy、Skimage等,可以写入requirements.txt,使用requirement插件一键安装。
numpy>=1.15.2
scikit-image>=0.14.0
torch
torchvision
pillow>=9.0.1
opencv-python>=4.5.5.62
glob2
scipy
matplotlib
1.5 算法代码
算法代码可以从github克隆或从 网盘(提取码:jack)下载 Jack提供的包体。github中没有训练数据,网盘包体可以直接运行。
2 算法运行
把你需要转换的图片(比如png)存放到./test_data/test_portrait_images/portrait_im/.目录下,然后在工程的主目录下敲下以下指令:
python u2net_portrait_test.py
你将在.\test_data\test_portrait_images\portrait_results中看到转换之后的肖像画效果。