最近由于项目需要做了一段时间的语义分割，希望能将自己的心路历程记录下来，以提供给所需帮助的人
接下来我将依托Unet语义分割网络介绍以下内容：
首先我的环境配置
pytorch1.10
win10
vs2017
python3.6
opencv3.4
Aaconda-5.2.0
一、使用pytorch实现简单的unet分割网络
二、使用Unet做多类别分割
三、c++调用python执行语义分割
四、c++调用libtorch执行语义分割
首先介绍使用pytorch实现简单的unet分割网络
Unet在医学图像分割领域是一个比较有名的分割网络，用Unet入坑语义分割由以下几个优点：
1、由于使用Unet模型做语义分割不需要大量的标注样本，甚至几十张标注样本都能取得较好的效果。
2、短时间内就能获取训练效果，可以根据自己的想法随意DIY Unet网络。
用pytorch实现Unet网络也很简单，下面是我复现的一个简单的Unet分割模型：
class conv_bn_relu(nn.Module):
def __init__(self, in_ch, out_ch,kerl = 3,pad = 1 ):
super(conv_bn_relu, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_ch, out_ch, kernel_size = kerl, padding = pad),
nn.BatchNorm2d(out_ch, eps=1e-3),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, kernel_size = kerl, padding = pad),
nn.BatchNorm2d(out_ch, eps=1e-3),
nn.ReLU(inplace=True)
)
def forward(self, input):
c = self.conv(input)
return c
class Unet(nn.Module):
def __init__(self,in_ch,out_ch):
super(Unet, self).__init__()
self.conv1 = conv_bn_relu(in_ch, 64)
self.pool1 = nn.MaxPool2d(2)
self.conv2 = conv_bn_relu(64, 128)
self.pool2 = nn.MaxPool2d(2)
self.conv3 = conv_bn_relu(128, 256)
self.up4 = nn.ConvTranspose2d(256, 128, 2, stride=2)
self.conv4 = conv_bn_relu(256, 128)
self.up5 = nn.ConvTranspose2d(128, 64, 2, stride=2)
self.conv5 = conv_bn_relu(128, 64)
self.conv6 = nn.Conv2d(64,out_ch, 1)
def forward(self,x):
c1=self.conv1(x)
p1=self.pool1(c1)
c2=self.conv2(p1)
p2=self.pool2(c2)
c3=self.conv3(p2)
up_4=self.up5(c3)
merge4 = torch.cat([up_4, c2], dim=1)
c4=self.conv4(merge4)
up_5=self.up4(c4)
merge5=torch.cat([up_5, c1],dim=1)
c5=self.conv5(merge5)
c6=self.conv6(c5)
return c6
显然Unet分割网络比其他语义分割模型更简洁，降低了小白入坑的门槛。
有了模型后我们如何去训练一个自己的模型呢，首先也是重要的一点就是数据读取，关于读取数据我写了一个简单的例子：
首先是获取样本图像的原始图像以及标注图像的路径
def make_dataset(root):
imgs=[]
path = os.path.join(root, "001.txt")
file = open(path)
for line in file.readlines():
line = line.split('\n')[0] #特别注意，去除文件路径后面的空白部分
img = os.path.join(root, "data", line)
mask = os.path.join(root, "label", line)
imgs.append((img,mask))
return imgs
我把原始数据与标注数据放在不同的文件夹下，且原始数据与标注数据的名称都相同。
所以文windows系统下只需要批处理生成001.txt文件就行了，具体操作如下：
首先在原始图像文件夹下建立一个test.txt文件，里面写上：
dir /b *.png *.jpg *.raw *.tif > 001.txt 这是相对路径
生成的只有图片文件名，形如 1.png
dir /s /b *.png *.jpg *.raw *.tif > 001.txt 这是绝对路径
生成的绝对路径，形如 D:\data\1.png
我们的txt文件中只需要写这一句 dir /b *.png *.jpg *.raw *.tif > 001.txt ，然后将test.txt文件名的后缀改为test.bat,然后双击运行就可以获得我们所需的txt文件
然后就是数据加载部分了，数据加载继承于Dataset，一般写法如下：
class LoadDataset(Dataset):
def __init__(self, root, transform=None, target_transform=None):
imgs = make_dataset(root)
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
def __getitem__(self, index):
x_path, y_path = self.imgs[index]
img_x = cv2.imread(x_path,0)
img_y = cv2.imread(y_path,0)
if self.transform is not None:
img_x = self.transform(img_x)
if self.target_transform is not None:
img_y = self.target_transform(img_y)
return img_x, img_y
def __len__(self):
return len(self.imgs)
接下来就轮到训练部分了
首先写上这两行，使得训练更稳定，具体原因可百度
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
原始数据与标注数据准换成Tensor
x_transforms = transforms.Compose([
transforms.ToTensor(),
#transforms.Normalize(mean=[0.85], std=[0.12]) #可根据需要决定
])
y_transforms = transforms.ToTensor()
训练模块如下：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def train_model(model, criterion, optimizer, dataload, num_epochs = 100):
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
dt_size = len(dataload.dataset)
epoch_loss = 0
step = 0
for x, y in dataload:
step += 1
inputs = x.to(device)
labels = y.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += float(loss.item())
print("%d/%d,train_loss:%0.5f" % (step, (dt_size - 1) // dataload.batch_size + 1, loss.item()))
torch.cuda.empty_cache() #及时清除无用缓存
print("epoch %d loss:%0.5f batch size " % (epoch, epoch_loss/step))
if np.isnan(epoch_loss/step):
break
torch.save(model.state_dict(), str(epoch)+ "_" + str(epoch_loss/step) + '_temp.pth') #保存每一轮的结果
torch.save(model.state_dict(), 'Unet.pth')
return model
def train():
model = Unet(1,1).to(device)
criterion = CrossEntropyLoss2d()
optimizer = optim.Adam(model.parameters(),lr=0.001)
liver_dataset = LoadDataset("D:/u-net/sample/train",transform = x_transforms,target_transform = y_transforms)
dataloaders = DataLoader(liver_dataset, batch_size = 4, shuffle=True, num_workers=4,pin_memory=True)
train_model(model, criterion, optimizer, dataloaders)
模型参数较大的情况下不推荐每轮保存一次模型，优化器只推荐Adam，batch_size视GPU显存的大小而定，pin_memory=True也需要根据GPU判定是否要开启（pytorch 默认为False），关于DataLoader参数设置可自行百度。
损失函数我们有更好的选择例如FocalLoss：
class FocalLoss(torch.nn.Module):
def __init__(self, gamma=2, alpha=0.25, reduction='elementwise_mean'):
super().__init__()
self.gamma = gamma
self.alpha = alpha
self.reduction = reduction
def forward(self, _input, target):
pt = torch.sigmoid(_input)
alpha = self.alpha
loss = - alpha * (1 - pt) ** self.gamma * target * torch.log(pt) - \
(1 - alpha) * pt ** self.gamma * (1 - target) * torch.log(1 - pt)
if self.reduction == 'elementwise_mean':
loss = torch.mean(loss)
elif self.reduction == 'sum':
loss = torch.sum(loss)
return loss
经测试二值语义分割FocalLoss要优于CrossEntropyLoss2d ，其他loss可自行尝试。
为了保证结果的可复现，我们每次训练手动初始化随机参数：
def seed_torch(seed=666):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
训练之前要测试以下数据处理是否正确：
def data_test():
liver_dataset = LoadDataset("D:/u-net/sample/train",transform=x_transforms,target_transform = y_transforms)
dataloaders = DataLoader(liver_dataset, batch_size=1, shuffle=True, num_workers=0)
for x, y in dataloaders:
img_x = np.uint8(torch.squeeze(x).numpy()*255)
img_y = np.uint8(torch.squeeze(y).numpy()*255)
print(img_x.shape)
print(img_y.shape)
cv2.imshow("img_x",img_x)
cv2.imshow("img_y",img_y)
cv2.waitKey(0)
break
另外需要注意：输入图像需要为2的倍数，或者保证所有层池化后反卷积能恢复同样的大小。
训练所需数据 可在kaggle下载
下一步篇将介绍如何用Unet做多类别分割