一、效果展示
（图片来源网络，如有侵权请联系删除）
对于天空和前景边界明显的场景，分割效果较好。如果边界不明显或者物体、衣物和场景天空类似，则可能效果不好。
如果期望达到更好的效果，除了对模型进一步训练外。还可以在获取到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实现天空分割，可用于替换天空变天效果
主业前端程序猿一枚。图片处理方面，作为业余爱好。如有错误，请各位大佬轻喷，谢谢！！😂
更多其他功能，可以扫下方二维码：
