pytorch+u2net实现天空分割

一、效果展示

（图片来源网络，如有侵权请联系删除）

对于天空和前景边界明显的场景，分割效果较好。如果边界不明显或者物体、衣物和场景天空类似，则可能效果不好。
如果期望达到更好的效果，除了对模型进一步训练外。还可以在获取到label图片后，进一步做处理。

二、摘要

用到的技术如下：

1、模型训练。使用pytorch版的u2net网络（https://gitee.com/throni0/U-2-Net）。

2、模型推理。使用onnxruntime，进行模型推理。模型训练中需要保存为onnx格式，或者训练完成后，将pytorch模型转换为onnx模型文件。

3、图像操作。使用Pillow，简单方便。

环境的配置，可自行查找资料。

三、模型训练

3.1 数据集

训练用的数据集，从网上找到了一些飞机分割的数据集，符合需求。又用ps做了一部分数据集。

将数据集放入项目中，运行u2net_train.py即可。

3.2 训练中断处理

一般训练需要很长时间，中途如果意外中断，或者希望调整参数，不得不停下。只需要稍微修改一下u2net_train.py即可。

# 载入模型结构
if model_name == 'u2net':
     net = U2NET(3, 1)
 elif model_name == 'u2netp':
     net = U2NETP(3, 1)

# 改为如下：
path_checkpoint = '之前已经训练好的模型文件'
resume = True   # 从头开始训练，将其设为False即可
if model_name == 'u2net':
    net = U2NET(3, 1)
    if resume:
        checkpoint = torch.load(path_checkpoint)
        net.load_state_dict(checkpoint)
elif model_name == 'u2netp':
    net = U2NETP(3, 1)

3.3 模型转换

将训练生成的pth文件，转换为onnx运行文件。

model = U2NET(3, 1)

pthfile = '训练生成的pth文件'
model.load_state_dict(torch.load(pthfile, map_location='cpu'))
model.eval()

img_input = torch.autograd.Variable(torch.randn(1, 3, 320, 320))   # 输入的图片类型，batch_size:1，图片张量纬度(3, 320, 320)
input_names = ["sky"]
output_names = ["output"]
torch.onnx.export(model, img_input, "./saved_models/sky/sky14.onnx", verbose=True, input_names=input_names, output_names=output_names, opset_version=11, training=False)

四、推理

推理使用onnxruntime

session = onnxruntime.InferenceSession("转换的onnx文件")
input_name = session.get_inputs()[0].name
label_name = session.get_outputs()[0].name

img_name_list = ['需要处理的图片']
image = Image.open(img_name_list[0])
w, h = image.size
dataset = SalObjDataset(
    img_name_list=img_name_list,
    lbl_name_list=[],
    transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)])
)
data_loader = DataLoader(
    dataset,
    batch_size=1,
    shuffle=False,
    num_workers=1
)
im = list(data_loader)[0]['image']
inputs_test = im
inputs_test = inputs_test.type(torch.FloatTensor)
with torch.no_grad():
    inputs_test = Variable(inputs_test)
res = session.run([label_name], {input_name: inputs_test.numpy().astype(np.float32)})
result = torch.from_numpy(res[0])
pred = result[:, 0, :, :]
pred = normPRED(pred)
pred = pred.squeeze()
predict_np = pred.cpu().data.numpy()
im = Image.fromarray(predict_np * 255).convert('RGB')
im = im.resize((w, h), resample=Image.BILINEAR)
im.show()

其他

这里给出资源，包括数据集、已经训练好的onnx模型、onnxruntime推理脚本，配置环境后可直接运行。如果你对结果不满意，也可以使用数据集，对模型继续进行训练。

资源路径：pytorch+u2net实现天空分割，可用于替换天空变天效果

第二版：pytorch+u2net实现天空分割（二）

主业前端程序猿一枚。图片处理方面，作为业余爱好。如有错误，请各位大佬轻喷，谢谢！！😂

更多其他功能，可以扫下方二维码：