pytorch+u2net实现天空分割

本文主要是介绍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].nameimg_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()