網絡結構

U-net是MICCAI2015的專門針對醫學影像分割設計的網絡結構，直到今天還是很好用，效果任然很好。網絡的左半部分是Encode
部分，此部分將輸入圖像逐步下采樣得到高層的語義信息。右半部分是Decoder部分，將Encode提取的語義特征逐步恢復為原始圖像的尺寸，中間是跳躍連接，將下采樣過程中提取到的不同level的特征加入到Decoder中，得到更好的分割結果。

完整訓練代碼和數據可以去我的github上面下載：https://github.com/SPECTRELWF/Semantic-Segmentation
個人主頁：liuweifeng.top:8090

數據集介紹

數據集是來自網上的視網膜血管分割數據集，訓練集只有二十張圖像，少得可憐。測試集也是二十張圖像。原始數據集中image和label使用.tif和.gif格式存儲的，需要先進行一下預處理，將數據組織成想要的樣子，再來進行訓練。

train

# !/usr/bin/python3 # -*- coding:utf-8 -*- # Author:WeiFeng Liu # @Time: 2021/12/9 下午1:54from unet_model.unet import UNET from dataset import SWXG_Dataset import torch.optim as optim import torch.nn as nn import torchdef train_net(net,device,data_path,epochs=40,batch_size=1,lr=1e-5):isbi_dataset = SWXG_Dataset(data_path)train_loader = torch.utils.data.DataLoader(isbi_dataset,batch_size,shuffle = True)#使用RMSprop優化optimizer = optim.RMSprop(net.parameters(),lr,weight_decay=1e-8,momentum=0.9)criterion = nn.BCEWithLogitsLoss()best_loss = float("inf")for epoch in range(epochs):net.train()for images, labels in train_loader:optimizer.zero_grad()images = images.to(device,dtype = torch.float32)labels = labels.to(device,dtype=torch.float32)pred = net(images)loss = criterion(pred,labels)print('epoch:%d train loss:%f' % (epoch+1,loss.item()))if loss <best_loss:best_loss = losstorch.save(net.state_dict(), 'best_model.pth')loss.backward()optimizer.step()if __name__ == "__main__":device = torch.device('cuda:0')net = UNET(n_channels=1,n_classes=1)net.to(device)data_path = r'Dataset/train/'train_net(net,device,data_path)

結果：

訓練30個epoch之后的結果看起來還是很好的。

總結

以上是生活随笔為你收集整理的Pytorch实现U-net视网膜血管分割的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。