Pytorch_Unet图像分割

GTX_AI

Pytorch_Unet图像分割：

1. # -*- coding: utf-8 -*-
   
2. import sys
   
3. sys.path.append(r"D:\2-LearningCode\999-AI-Pytorch\3_AI_nets\u_net_liver-master\u_net_liver-master")
   
4. import numpy as np
   
5. import torch
   
6. import argparse
   
7. from torch.utils.data import DataLoader
   
8. from torch import autograd, optim
   
9. from torchvision.transforms import transforms
   
10. from unet import Unet
    
11. from dataset import LiverDataset
    
12. import cv2
    
13. import random
    
14. 
    
15. # 是否使用cuda
    
16. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
17. #device = torch.device("cpu")
    
18. #print(device)
    
19. 
    
20. x_transforms = transforms.Compose([
    
21.     transforms.ToTensor(),
    
22.     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    
23. ])
    
24. 
    
25. # mask只需要转换为tensor
    
26. y_transforms = transforms.ToTensor()
    
27. 
    
28. #参数解析
    
29. parse=argparse.ArgumentParser()
    
30. 
    
31. def train_model(model, criterion, optimizer, dataload, num_epochs=20):
    
32.     model.cuda(0)
    
33.     for epoch in range(num_epochs):
    
34.         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
    
35.         print('-' * 10)
    
36.         dt_size = len(dataload.dataset)
    
37.         epoch_loss = 0
    
38.         step = 0
    
39.         #with torch.no_grad():
    
40.         for x, y in dataload:
    
41.             step += 1
    
42.             inputs = x.to(device)
    
43.             labels = y.to(device)
    
44.             del x,y
    
45.             # zero the parameter gradients
    
46.             optimizer.zero_grad()
    
47.             # forward
    
48.             outputs = model(inputs)
    
49.             loss = criterion(outputs, labels)
    
50.             loss.backward()
    
51.             optimizer.step()
    
52.             
    
53.             epoch_loss += loss.item()
    
54.             print("%d/%d,train_loss:%0.3f" % (step, (dt_size - 1) // dataload.batch_size + 1, loss.item()))
    
55.             
    
56.             del inputs,labels,outputs
    
57.             torch.cuda.empty_cache()
    
58.             
    
59.         print("epoch %d loss:%0.3f" % (epoch, epoch_loss))
    
60.     torch.save(model.state_dict(), 'weights_%d.pth' % epoch)
    
61.     return model
    
62. 
    
63. #训练模型
    
64. def train():
    
65.     model = Unet(3, 1).to(device)
    
66.     batch_size = args.batch_size
    
67.     criterion = torch.nn.BCELoss()
    
68.     optimizer = optim.Adam(model.parameters())
    
69.     liver_dataset = LiverDataset("data/train",transform=x_transforms,target_transform=y_transforms)
    
70.     dataloaders = DataLoader(liver_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    
71.     train_model(model, criterion, optimizer, dataloaders)
    
72. 
    
73. #显示模型的输出结果
    
74. def test_gpu():
    
75.     model = Unet(3, 1).to(device)
    
76. #    model.load_state_dict(torch.load(args.ckp,map_location='cpu'))
    
77.     model.load_state_dict(torch.load(args.ckp, map_location='cuda'))
    
78.     liver_dataset = LiverDataset("data/val", transform=x_transforms,target_transform=y_transforms)
    
79.     dataloaders = DataLoader(liver_dataset, batch_size=1)
    
80.     model.cuda(0)
    
81.     model.eval()
    
82.     import matplotlib.pyplot as plt
    
83.     plt.ion()
    
84.     with torch.no_grad():
    
85.         for x, _ in dataloaders:
    
86.             y=model(x.to(device))
    
87.             img_y=torch.squeeze(y).cpu().numpy()
    
88.             plt.imshow(img_y)
    
89.             plt.pause(0.01)
    
90.             
    
91.             torch.cuda.empty_cache()
    
92.             
    
93.             cv2.imwrite('./resultsImages/'+str(random.random())+'.tiff', img_y)
    
94. 
    
95.         plt.show()
    
96.     torch.cuda.empty_cache()
    
97.    
    
98. def test_cpu():
    
99.     model = Unet(3, 1)
    
100.     model.load_state_dict(torch.load(args.ckp,map_location='cpu'))
     
101.     liver_dataset = LiverDataset("data/val", transform=x_transforms,target_transform=y_transforms)
     
102.     dataloaders = DataLoader(liver_dataset, batch_size=1)
     
103.     model.eval()
     
104.     import matplotlib.pyplot as plt
     
105.     plt.ion()
     
106.     with torch.no_grad():
     
107.         for x, _ in dataloaders:
     
108.             y=model(x)
     
109.             img_y=torch.squeeze(y).numpy()
     
110.             plt.imshow(img_y)
     
111.             plt.pause(0.01)
     
112.         plt.show()
     
113.         
     
114.     torch.cuda.empty_cache()
     
115.    
     
116. if __name__ == '__main__':
     
117.     parse = argparse.ArgumentParser()
     
118. #    parse.add_argument("action", type=str, help="train or test")
     
119.     parse.add_argument("--action", type=str, help="train or test", default="train")
     
120.     parse.add_argument("--batch_size", type=int, default=4)
     
121. #    parse.add_argument("--ckp", type=str, help="the path of model weight file")
     
122.     parse.add_argument("--ckp", type=str, help="the path of model weight file", default="weights_19.pth")
     
123.     args = parse.parse_args()
     
124. 
     
125.     print(args.action)
     
126. 
     
127.     if args.action=="train":
     
128.         train()
     
129.     elif args.action=="test":
     
130.         test_gpu()
     
131.         
     
132.     test_gpu()
     
133.    

复制代码

尤其要注意的是：
del inputs,labels,outputs
torch.cuda.empty_cache()
清理GPU缓存，不然报GPU内存不够错误。

unet.py如下：

1. import torch.nn as nn
   
2. import torch
   
3. from torch import autograd
   
4. 
   
5. class DoubleConv(nn.Module):
   
6.     def __init__(self, in_ch, out_ch):
   
7.         super(DoubleConv, self).__init__()
   
8.         self.conv = nn.Sequential(
   
9.             nn.Conv2d(in_ch, out_ch, 3, padding=1),
   
10.             nn.BatchNorm2d(out_ch),
    
11.             nn.ReLU(inplace=True),
    
12.             nn.Conv2d(out_ch, out_ch, 3, padding=1),
    
13.             nn.BatchNorm2d(out_ch),
    
14.             nn.ReLU(inplace=True)
    
15.         )
    
16. 
    
17.     def forward(self, input):
    
18.         return self.conv(input)
    
19. 
    
20. 
    
21. class Unet(nn.Module):
    
22.     def __init__(self,in_ch,out_ch):
    
23.         super(Unet, self).__init__()
    
24. 
    
25.         self.conv1 = DoubleConv(in_ch, 64)
    
26.         self.pool1 = nn.MaxPool2d(2)
    
27.         self.conv2 = DoubleConv(64, 128)
    
28.         self.pool2 = nn.MaxPool2d(2)
    
29.         self.conv3 = DoubleConv(128, 256)
    
30.         self.pool3 = nn.MaxPool2d(2)
    
31.         self.conv4 = DoubleConv(256, 512)
    
32.         self.pool4 = nn.MaxPool2d(2)
    
33.         self.conv5 = DoubleConv(512, 1024)
    
34.         self.up6 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
    
35.         self.conv6 = DoubleConv(1024, 512)
    
36.         self.up7 = nn.ConvTranspose2d(512, 256, 2, stride=2)
    
37.         self.conv7 = DoubleConv(512, 256)
    
38.         self.up8 = nn.ConvTranspose2d(256, 128, 2, stride=2)
    
39.         self.conv8 = DoubleConv(256, 128)
    
40.         self.up9 = nn.ConvTranspose2d(128, 64, 2, stride=2)
    
41.         self.conv9 = DoubleConv(128, 64)
    
42.         self.conv10 = nn.Conv2d(64,out_ch, 1)
    
43. 
    
44.     def forward(self,x):
    
45.         c1=self.conv1(x)
    
46.         p1=self.pool1(c1)
    
47.         c2=self.conv2(p1)
    
48.         p2=self.pool2(c2)
    
49.         c3=self.conv3(p2)
    
50.         p3=self.pool3(c3)
    
51.         c4=self.conv4(p3)
    
52.         p4=self.pool4(c4)
    
53.         c5=self.conv5(p4)
    
54.         up_6= self.up6(c5)
    
55.         merge6 = torch.cat([up_6, c4], dim=1)
    
56.         c6=self.conv6(merge6)
    
57.         up_7=self.up7(c6)
    
58.         merge7 = torch.cat([up_7, c3], dim=1)
    
59.         c7=self.conv7(merge7)
    
60.         up_8=self.up8(c7)
    
61.         merge8 = torch.cat([up_8, c2], dim=1)
    
62.         c8=self.conv8(merge8)
    
63.         up_9=self.up9(c8)
    
64.         merge9=torch.cat([up_9,c1],dim=1)
    
65.         c9=self.conv9(merge9)
    
66.         c10=self.conv10(c9)
    
67.         out = nn.Sigmoid()(c10)
    
68.         return out

复制代码

参考：
【1】读取文件夹内图像--方法1
【2】读取文件夹内图像--方法2
【3】https://github.com/gupta-abhay/pytorch-modelzoo
【4】https://github.com/yt4766269/pytorch_zoo
【5】https://github.com/takahiro-itazuri/model-zoo-pytorch
【6】http://222.195.93.137/gitlab/winston.wen/kaggle-1
【7】3-6之百度网盘下载链接：https://pan.baidu.com/s/17pXf2M3lIAFkJuoMXOZ9pA 提取码：k0ou