-
Notifications
You must be signed in to change notification settings - Fork 38
卷积模型人像分割
larry edited this page Jul 26, 2020
·
3 revisions
人像分割
我们拍摄照片一般没有条件准备蓝底和红底背景布,需要通过技术手段将人像和背景进行分离,然后加入背景,
这里的人像分割采用的是 u2net
u2net :https://github.com/NathanUA/U-2-Net
分割效果:
安装u2net 依赖包
安装pytorch
进入你的python环境,输入 pytorch 官网提供的命令
(idphoto) C:\Users\larry>conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
安装 scikit-image
方法1:如果安装了git,直接clone 到目录就行
git clone https://github.com/NathanUA/U-2-Net.git
方法2:可以直接在网页上打开url访问,然后点击code 按钮,下载代码
下载解压之后的目录结构
使用之前需要下载训练好的模型 u2net.pth
这个地址是在google服务器上,有些网友可能没法访问,可以联系我私发
这边改造了一个方法可以传递一个图片,然后转换成 trimap
def test_seg_trimap(org,org_trimap,resize_trimap):
# 将原始图片转换成trimap
# org:原始图片
# org_trimap:
# resize_trimap: 调整尺寸的trimap
image = Image.open(org)
print(image)
img = np.array(image)
net = pre_net()
inputs_test = pre_test_data(img)
d1, d2, d3, d4, d5, d6, d7 = net(inputs_test)
# normalization
pred = d1[:, 0, :, :]
pred = normPRED(pred)
# 将数据转换成图片
im = get_im(pred)
im.save(org_trimap)
sp = image.size
# 根据原始图片调整尺寸
imo = im.resize((sp[0], sp[1]), resample=Image.BILINEAR)
imo.save(resize_trimap)