这是一个开源项目，我已经成功在集群上部署，并且已经跑了一遍。

整个项目可以有两个效果

修复图片的缺失
去掉不想要的元素

效果还是蛮酷炫的。

参考内容

安装过程

安装背景

ubuntu16.04
anaconda

创建虚拟环境

conda create -n edge python=3.5

并且进入虚拟环境

clone 代码

选择自己喜欢的文件夹

git clone https://github.com/knazeri/edge-connect.git
cd edge-connect

安装 torch

pip install torch==1.0.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install torchvision==0.2.1 -i https://pypi.tuna.tsinghua.edu.cn/simple

安装其它依赖

pip install -r requirements.txt

下载官方的 checkpoints

checkponits/r8v4

假设，你的项目地址是

/data/aaa/edge-connect

然后移动 checkpoint

cd /data/aaa/edge-connect
mkdir checkpoints
mv path_of_master.zip /data/aaa/edge-connect/checkpoints

这个 master.zip 就是上面打包的 checkpoints 压缩包。

cd /data/aaa/edge-connect/checkpoints
unzip master.zip
这个时候可以 ls 看一下

这个时候，基本的配置都已经做好了。

运行

python test.py \
  --model [stage] \
  --checkpoints [path to checkpoints] \
  --input [path to input directory or file] \
  --mask [path to masks directory or mask file] \
  --output [path to the output directory]

这个要解释一下

model :这个 model 默认是 3，其他的我没有试过，可以自己读源码进行尝试
checkpoints 这个是官方已经训练好的参数集，一共有 3 个，我猜测，可能和 model 一一对应，所以，这点要注意
input 这个是要处理的图片文件夹
mask 这个是处理图片中每一个图片使用的遮罩
output 这个是要输出的结果图片的文件夹

图片和遮罩

我们是预处理图片和遮罩一一对应。

这里面的一一对应有两个含义

名字相同，预处理图片和遮罩图片名字相同
遮罩路径和预处理图片中的路径相同
对应图片长宽必须相等

举一个例子，假设我们的预处理图片放在叫做

/data/aaa/edge-connect/example/places/images

里面，并且叫做

places2_01.png

那么，同样在

/data/aaa/edge-connect/example/places/masks

里面需要有一个遮罩叫做

places2_01.png

上面的路径是我胡乱编的，但是，其核心点在于名字要相同。

那我们来看看这两个文件究竟有什么相似之处。

预处理图片

RBG 模式
最好是 256 * 256
通道显示为（w,h,3）

遮罩图片

gray 模式
最好是 256 * 256
alpha 通道：否

这里要着重强调，预测图和遮罩图并不一定是 256 * 256 ，但是，必须长宽相等。

可以看出其路径是一样的，要特别注意，路径必须相同，否则会出现不可预知的情况。

这里，我将提供一个程序，是将遮罩和原图片进行混合的。

import cv2
import numpy as np

a = cv2.imread('./places2_02.png')
b = cv2.imread('./places2_01.png')

cv2.imwrite('./test.png', np.bitwise_or(b, a))

这里面 places2_01.png 代指遮罩，也就是

/places2_02.png) 是指原图片，也就是

经过合成，就形成了

为了让流程走的更顺利，我创建了一个叫做 init_png.py 的文件，其内容如下:

import cv2
import glob
import argparse
import numpy as np


def init_png(origin_dir, mask_dir, result_dir):
    origin_pngs = glob.glob(f'{origin_dir}/*.png')
    for origin_png in origin_pngs:
        a = cv2.imread(origin_png)
        a = cv2.resize(a, (256, 256))
        b = cv2.imread(f'{mask_dir}/{origin_png.split("/")[-1]}')
        cv2.imwrite(f'{result_dir}/{origin_png.split("/")[-1]}', np.bitwise_or(b, a))
        print(f'图片 {origin_png.split("/")[-1]} 生成成功')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--origin_dir")
    parser.add_argument("--mask_dir")
    parser.add_argument("--result_dir")
    args = parser.parse_args()
    init_png(args.origin_dir, args.mask_dir, args.result_dir)

其中

origin_dir:是原始图像的文件夹
mask_dir：是遮罩图片的文件夹
result_dir：是合成图片的文件夹

所以，你可以先初始化对象。

运行代码

举一个简单的运行代码的例子

python test.py \
  --checkpoints /data/aaa/edge-connect/checkpoints/places2  \
  --input /data/aaa/edge-connect/examples/places2/images \
  --mask /data/aaa/edge-connect/examples/places2/masks \
  --output /data/aaa/edge-connect/checkpoints/resluts

其中，output 的文件夹不需要自己建。

最后运行会非常流畅。

其他

出现异常

或许和你图片的权限有关，你可以使用

ls -lah

看看你是否有这些图片的权限，如果没有，可以进入管理员权限，运行

chmod 777 文件名