Vgg16预训练下的Unet

GTX_AI

Vgg16预训练下的Unet

1. # -*- coding: utf-8 -*-
   
2. """
   
3. Created on Tue Aug 27 22:06:49 2019
   
4. 
   
5. @author: Solem
   
6. """
   
7. import sys
   
8. sys.path.append(r"D:\2-LearningCode\999-AI-Pytorch\3_AI_nets\u_net_liver-master\Vgg16_Unet")
   
9. import os
   
10. os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    
11. import numpy as np
    
12. import torch
    
13. import torch.utils.data as data
    
14. from torch import autograd, optim
    
15. from torchvision.transforms import transforms
    
16. from torch.autograd import Variable
    
17. # 配置文件
    
18. from configs import configXML
    
19. from pytorch_zoo import unet, resnet38unet
    
20. # begin
    
21. net = unet.Vgg16bn(num_classes=configXML.num_classes)
    
22. net.train()
    
23. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
24. #device = torch.device("cpu")
    
25. #print(device)
    
26. net = net.cuda()
    
27. criterion1 =  torch.nn.BCELoss()
    
28. if configXML.initial_checkpoints is not None:
    
29.     net = torch.load(configXML.initial_checkpoints)
    
30. best_val_loss = np.inf
    
31. num_epoch_has_trained = 0
    
32. for epoch in range(num_epoch_has_trained, configXML.num_epochs):
    
33.     iter_count = 0;
    
34.     train_losses = 0.0;
    
35.     optimizer = optim.SGD(net.parameters(), lr = configXML.lr_init, momentum = configXML.moment_init)
    
36.     torch.cuda.empty_cache()
    
37.     net.train()
    
38.     for i,(train_image,train_mask) in enumerate(configXML.train_loader):
    
39.         iter_count = iter_count+1
    
40.         train_image = Variable(train_image.cuda())
    
41.         train_mask = Variable(train_mask.cuda())
    
42.         train_logits = net(train_image)
    
43.         train_loss = criterion1(train_logits, train_mask)
    
44.         train_losses = train_losses + train_loss.item()
    
45.         optimizer.zero_grad()
    
46.         train_loss.backward()
    
47.         optimizer.step()
    
48.     torch.cuda.empty_cache()   
    
49.     torch.save(net.state_dict(), './checkPoints/Vgg16_Val_normal.pth')

复制代码

configXML.py:

1. # -*- coding: utf-8 -*-
   
2. """
   
3. Created on Sun Sep  8 20:47:37 2019
   
4. 
   
5. @author: Solem
   
6. """
   
7. from XMLdataSet import dataset_improved, dataset_package
   
8. import torch
   
9. from torchvision.transforms import transforms
   
10. from torch.utils.data import DataLoader
    
11. 
    
12. Train_file=r'D:\2-LearningCode\992-DataEnhanced\Samples\EnhancedImages\Images_train.txt'
    
13. Val_file=r'D:\2-LearningCode\992-DataEnhanced\Samples\EnhancedImages\Images_val.txt'
    
14. Imagepath=r"D:\2-LearningCode\992-DataEnhanced\Samples\EnhancedImages\Images"
    
15. Images_file=r'D:\2-LearningCode\992-DataEnhanced\Samples\EnhancedImages\Images.txt'
    
16. Maskspath=r"D:\2-LearningCode\992-DataEnhanced\Samples\EnhancedImages\Masks"
    
17. Masks_file=r'D:\2-LearningCode\992-DataEnhanced\Samples\EnhancedImages\Masks.txt'
    
18. num_classes = 1
    
19. img_size = 256
    
20. batch_size = 4
    
21. num_workers = 0
    
22. 
    
23. initial_checkpoints = None
    
24. #initial_checkpoints = r'D:\2-LearningCode\999-AI-Pytorch\3_AI_nets\u_net_liver-master\Vgg16_Unet\checkPoints\Vgg16_Val_normal.pth'
    
25. num_epochs = 1000
    
26. lr_init = 0.01
    
27. moment_init = 0.9
    
28. 
    
29. # 均值0.5，方差0.5
    
30. x_transforms = transforms.Compose([
    
31.     transforms.ToTensor(),
    
32.     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    
33. ])
    
34. 
    
35. # mask只需要转换为tensor
    
36. y_transforms = transforms.ToTensor()
    
37. 
    
38. #train_liver_dataset = dataset_improved.LiverDataset(imagepath, Train_file, transform=x_transforms, target_transform=y_transforms)
    
39. #train_loader = DataLoader(train_liver_dataset, batch_size=batch_size, shuffle=True, num_workers = num_workers)
    
40. 
    
41. train_dataset = dataset_package.opencvDateset(img_root_path=Imagepath, mask_root_path=Maskspath, txt_file=Train_file)
    
42. train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    
43. 
    
44. #val_liver_dataset = dataset_improved.LiverDataset(imagepath, Val_file, transform=x_transforms, target_transform=y_transforms)
    
45. #val_loader = DataLoader(val_liver_dataset, batch_size=batch_size, shuffle=True, num_workers = num_workers)
    
46. 
    
47. val_dataset = dataset_package.opencvDateset(img_root_path=Imagepath, mask_root_path=Maskspath, txt_file=Val_file)
    
48. val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)

复制代码

dataset_package.py如下：

1. # -*- coding: utf-8 -*-
   
2. """
   
3. Created on Sun May 19 22:28:08 2019
   
4. 
   
5. @author: Solem
   
6. """
   
7. import os
   
8. import PIL.Image as Image
   
9. import numpy as np
   
10. import cv2
    
11. import torch
    
12. import torch.utils.data as data
    
13. from torch import autograd, optim
    
14. from torchvision.transforms import transforms
    
15. from torch.autograd import Variable
    
16. from configs import configXML
    
17. 
    
18. def make_dataset(root):
    
19.     imgs=[]
    
20. #    n=len(os.listdir(root))//2
    
21. #    for i in range(100):
    
22.     for i in range(100):
    
23.         for j in range(16):
    
24.             img=os.path.join(root,"%03d_%d.png"%(i,j+1))
    
25.             mask=os.path.join(root,"%03d_%d_mask.png"%(i,j+1))
    
26.             imgs.append((img,mask))
    
27.     return imgs
    
28. 
    
29. class LiverDataset(data.Dataset):
    
30.     def __init__(self, root, transform=None, target_transform=None):
    
31.         imgs = make_dataset(root)
    
32.         self.imgs = imgs
    
33.         self.transform = transform
    
34.         self.target_transform = target_transform
    
35. 
    
36.     def __getitem__(self, index):
    
37.         x_path, y_path = self.imgs[index]
    
38.         img_x = Image.open(x_path)
    
39.         img_y = Image.open(y_path)
    
40.         if self.transform is not None:
    
41.             img_x = self.transform(img_x)
    
42.         if self.target_transform is not None:
    
43.             img_y = self.target_transform(img_y)
    
44.         return img_x, img_y
    
45. 
    
46.     def __len__(self):
    
47.         return len(self.imgs)
    
48. 
    
49. class opencvDateset(data.Dataset):
    
50.     def __init__(self,img_root_path,mask_root_path, txt_file, transform=None):
    
51.         self.transform = transform
    
52.         self.img_root_path = img_root_path
    
53.         self.mask_root_path = mask_root_path
    
54.         with open(txt_file) as f:
    
55.             self.indexs = f.readlines()
    
56.     def __getitem__(self, idx):      
    
57.         index = self.indexs[idx].split('.')[0]
    
58.         image_path = os.path.join(self.img_root_path, index+'.jpg')
    
59.         mask_path = os.path.join(self.mask_root_path, index+'_label.png')
    
60.         img = cv2.imread(image_path)
    
61.         img = cv2.resize(img, (configXML.img_size, configXML.img_size))
    
62.         img = img.astype(np.float32)/255.0
    
63.         mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    
64. #        print(np.shape(mask))
    
65.         mask = cv2.resize(mask, (configXML.img_size, configXML.img_size))
    
66.         print(np.shape(mask))
    
67.         if self.transform is not None:
    
68.             for trans in self.transform:
    
69.                 img, mask = trans(img, mask)
    
70.         img = torch.FloatTensor(img)
    
71.         mask = torch.FloatTensor(mask)
    
72.         return img, mask
    
73.    
    
74.     def __len__(self):
    
75.         return len(self.indexs)

复制代码

参考：
【1】Pytorch_Unet图像分割
【2】http://222.195.93.137/gitlab/winston.wen/kaggle-1

Halcom

dataset_package.py的第70行和第71行，应该可以这样修改，不然容易报错：

1. def image_to_tensor(image, mean = 0.0, std = 1.0):
   
2.         image = (image-mean)/std
   
3.         image = image.transpose(2,0,1)
   
4.         image = torch.from_numpy(image)
   
5.         return image
   
6. 
   
7. def mask_to_tensor(mask):
   
8.         mask = (mask>128).astype(np.float32)
   
9.         mask = torch.from_numpy(mask)
   
10.         return mask

复制代码

第61行和第65行，是否显得多余。

针对主程序而言：
from configs import configXML
这样import，那么你确实需要configXML.initial_checkpoints，configXML.num_epochs来调用类的成员变量
可以改为
from configs.configXML import *
然后你使用的都是全局变量了，可以直接使用initial_checkpoints，num_epochs
记得每个文件夹放一个__init__.py文件。