Skip to content

卷积模型人像分割

larry edited this page Jul 26, 2020 · 3 revisions

人像分割

我们拍摄照片一般没有条件准备蓝底和红底背景布,需要通过技术手段将人像和背景进行分离,然后加入背景,

这里的人像分割采用的是 u2net

u2net :https://github.com/NathanUA/U-2-Net

分割效果:

image-20200704182505300

image-20200704182533026

安装u2net 依赖包

安装pytorch

进入你的python环境,输入 pytorch 官网提供的命令

image-20200602163429494

(idphoto) C:\Users\larry>conda install pytorch torchvision cudatoolkit=10.2 -c pytorch

image-20200602161424515

安装 scikit-image

image-20200602163741238

image-20200602171726689

image-20200604172218818

使用u2net

方法1:如果安装了git,直接clone 到目录就行

git clone https://github.com/NathanUA/U-2-Net.git

方法2:可以直接在网页上打开url访问,然后点击code 按钮,下载代码

image-20200704112548211

下载解压之后的目录结构

image-20200704114151461

使用之前需要下载训练好的模型 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)