实战使用

戴琼海老师组3D-Unet无监督去噪

data_process.py

  1. import numpy as np
  2. import os
  3. import tifffile as tiff
  4. import random
  5. import math
  6. import torch
  7. from torch.utils.data import Dataset
  9. class trainset(Dataset):
  10.     def __init__(self,name_list,coordinate_list,noise_img_all,stack_index):
  11.         self.name_list = name_list
  12.         self.coordinate_list=coordinate_list
  13.         self.noise_img_all = noise_img_all
  14.         self.stack_index = stack_index
  16.     def __getitem__(self, index):
  17.         #fn = self.images[index]
  18.         stack_index = self.stack_index[index]
  19.         noise_img = self.noise_img_all[stack_index]
  20.         single_coordinate = self.coordinate_list[self.name_list[index]]
  21.         init_h = single_coordinate['init_h']
  22.         end_h = single_coordinate['end_h']
  23.         init_w = single_coordinate['init_w']
  24.         end_w = single_coordinate['end_w']
  25.         init_s = single_coordinate['init_s']
  26.         end_s = single_coordinate['end_s']
  27.         input = noise_img[init_s:end_s:2, init_h:end_h, init_w:end_w]
  28.         target = noise_img[init_s + 1:end_s:2, init_h:end_h, init_w:end_w]
  30.         input=torch.from_numpy(np.expand_dims(input, 0))
  31.         target=torch.from_numpy(np.expand_dims(target, 0))
  32.         return input, target
  34.     def __len__(self):
  35.         return len(self.name_list)
  38. class testset(Dataset):
  39.     def __init__(self,name_list,coordinate_list,noise_img):
  40.         self.name_list = name_list
  41.         self.coordinate_list=coordinate_list
  42.         self.noise_img = noise_img
  44.     def __getitem__(self, index):
  45.         #fn = self.images[index]
  46.         single_coordinate = self.coordinate_list[self.name_list[index]]
  47.         init_h = single_coordinate['init_h']
  48.         end_h = single_coordinate['end_h']
  49.         init_w = single_coordinate['init_w']
  50.         end_w = single_coordinate['end_w']
  51.         init_s = single_coordinate['init_s']
  52.         end_s = single_coordinate['end_s']
  53.         noise_patch = self.noise_img[init_s:end_s, init_h:end_h, init_w:end_w]
  54.         noise_patch=torch.from_numpy(np.expand_dims(noise_patch, 0))
  55.         #target = self.target[index]
  56.         return noise_patch,single_coordinate
  58.     def __len__(self):
  59.         return len(self.name_list)
  61. def train_preprocess_lessMemoryMulStacks(args):
  62.     img_h = args.img_h
  63.     img_w = args.img_w
  64.     img_s2 = args.img_s*2
  65.     gap_h = args.gap_h
  66.     gap_w = args.gap_w
  67.     gap_s2 = args.gap_s*2
  68.     im_folder = args.datasets_path + '//' + args.datasets_folder
  70.     name_list = []
  71.     coordinate_list={}
  72.     stack_index = []
  73.     noise_im_all = []
  74.     ind = 0;
  75.     print('\033[1;31mImage list for training -----> \033[0m')
  76.     stack_num = len(list(os.walk(im_folder, topdown=False))[-1][-1])
  77.     print('Total number -----> ', stack_num)
  78.     for im_name in list(os.walk(im_folder, topdown=False))[-1][-1]:
  79.         print(im_name)
  80.         im_dir = im_folder+ '//' + im_name
  81.         noise_im = tiff.imread(im_dir)
  82.         if noise_im.shape[0]>args.select_img_num:
  83.             noise_im = noise_im[0:args.select_img_num,:,:]
  84.         gap_s2 = get_gap_s(args, noise_im, stack_num)
  85.         # print('noise_im shape -----> ',noise_im.shape)
  86.         # print('noise_im max -----> ',noise_im.max())
  87.         # print('noise_im min -----> ',noise_im.min())
  88.         noise_im = (noise_im-noise_im.min()).astype(np.float32)/args.normalize_factor 
  89.         noise_im_all.append(noise_im)
  91.         whole_w = noise_im.shape[2]
  92.         whole_h = noise_im.shape[1]
  93.         whole_s = noise_im.shape[0]
  94.         # print('int((whole_h-img_h+gap_h)/gap_h) -----> ',int((whole_h-img_h+gap_h)/gap_h))
  95.         # print('int((whole_w-img_w+gap_w)/gap_w) -----> ',int((whole_w-img_w+gap_w)/gap_w))
  96.         # print('int((whole_s-img_s2+gap_s2)/gap_s2) -----> ',int((whole_s-img_s2+gap_s2)/gap_s2))
  97.         for x in range(0,int((whole_h-img_h+gap_h)/gap_h)):
  98.             for y in range(0,int((whole_w-img_w+gap_w)/gap_w)):
  99.                 for z in range(0,int((whole_s-img_s2+gap_s2)/gap_s2)):
  100.                     single_coordinate={'init_h':0, 'end_h':0, 'init_w':0, 'end_w':0, 'init_s':0, 'end_s':0}
  101.                     init_h = gap_h*x
  102.                     end_h = gap_h*x + img_h
  103.                     init_w = gap_w*y
  104.                     end_w = gap_w*y + img_w
  105.                     init_s = gap_s2*z
  106.                     end_s = gap_s2*z + img_s2
  107.                     single_coordinate['init_h'] = init_h
  108.                     single_coordinate['end_h'] = end_h
  109.                     single_coordinate['init_w'] = init_w
  110.                     single_coordinate['end_w'] = end_w
  111.                     single_coordinate['init_s'] = init_s
  112.                     single_coordinate['end_s'] = end_s
  113.                     # noise_patch1 = noise_im[init_s:end_s,init_h:end_h,init_w:end_w]
  114.                     patch_name = args.datasets_folder+'_'+im_name.replace('.tif','')+'_x'+str(x)+'_y'+str(y)+'_z'+str(z)
  115.                     # train_raw.append(noise_patch1.transpose(1,2,0))
  116.                     name_list.append(patch_name)
  117.                     # print(' single_coordinate -----> ',single_coordinate)
  118.                     coordinate_list[patch_name] = single_coordinate
  119.                     stack_index.append(ind)
  120.         ind = ind + 1;
  121.     return  name_list, noise_im_all, coordinate_list, stack_index

train.py

实例化trainset类train_data
DataLoader

  1. # start training
  2. for epoch in range(0, opt.n_epochs):
  3.     train_data = trainset(name_list, coordinate_list, noise_img_all,stack_index)
  4.     trainloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=4)
  5.     for iteration, (input, target) in enumerate(trainloader):
  6.         input=input.cuda()
  7.         target = target.cuda()
  8.         real_A=input
  9.         real_B=target
  10.         real_A = Variable(real_A)
  11.         #print('real_A shape -----> ', real_A.shape)
  12.         #print('real_B shape -----> ',real_B.shape)
  13.         fake_B = denoise_generator(real_A)
  14.         L1_loss = L1_pixelwise(fake_B, real_B)
  15.         L2_loss = L2_pixelwise(fake_B, real_B)
  1. import os
  2. import torch
  3. import torch.nn as nn
  4. from torch.autograd import Variable
  5. from torch.utils.data import DataLoader
  6. import argparse
  7. import time
  8. import datetime
  9. from network import Network_3D_Unet
  10. from data_process import train_preprocess_lessMemoryMulStacks, trainset
  11. from utils import save_yaml
  13. #############################################################################################################################################
  14. parser = argparse.ArgumentParser()
  15. parser.add_argument("--n_epochs", type=int, default=40, help="number of training epochs")
  16. parser.add_argument('--GPU', type=str, default='0,1', help="the index of GPU you will use for computation")
  18. parser.add_argument('--batch_size', type=int, default=2, help="batch size")
  19. parser.add_argument('--img_w', type=int, default=150, help="the width of image patch")
  20. parser.add_argument('--img_h', type=int, default=150, help="the height of image patch")
  21. parser.add_argument('--img_s', type=int, default=150, help="the length of image patch")
  23. parser.add_argument('--lr', type=float, default=0.00005, help='initial learning rate')
  24. parser.add_argument("--b1", type=float, default=0.5, help="Adam: bata1")
  25. parser.add_argument("--b2", type=float, default=0.999, help="Adam: bata2")
  26. parser.add_argument('--normalize_factor', type=int, default=1, help='normalize factor')
  27. parser.add_argument('--fmap', type=int, default=16, help='number of feature maps')
  29. parser.add_argument('--output_dir', type=str, default='./results', help="output directory")
  30. parser.add_argument('--datasets_folder', type=str, default='train', help="A folder containing files for training")
  31. parser.add_argument('--datasets_path', type=str, default='datasets', help="dataset root path")
  32. parser.add_argument('--pth_path', type=str, default='pth', help="pth file root path")
  33. parser.add_argument('--select_img_num', type=int, default=100000, help='select the number of images used for training')
  34. parser.add_argument('--train_datasets_size', type=int, default=4000, help='datasets size for training')
  35. opt = parser.parse_args()
  37. # default image gap is 0.5*image_dim
  38. # opt.gap_s (image gap) is the distance between two adjacent patches
  39. opt.gap_s=int(opt.img_s*0.5)
  40. opt.gap_w=int(opt.img_w*0.5)
  41. opt.gap_h=int(opt.img_h*0.5)
  42. opt.ngpu=str(opt.GPU).count(',')+1
  43. print('\033[1;31mTraining parameters -----> \033[0m')
  44. print(opt)
  46. ########################################################################################################################
  47. if not os.path.exists(opt.output_dir): 
  48.     os.mkdir(opt.output_dir)
  49. current_time = opt.datasets_folder+'_'+datetime.datetime.now().strftime("%Y%m%d%H%M")
  50. output_path = opt.output_dir + '/' + current_time
  51. pth_path = 'pth//'+ current_time
  52. if not os.path.exists(pth_path): 
  53.     os.mkdir(pth_path)
  55. yaml_name = pth_path+'//para.yaml'
  56. save_yaml(opt, yaml_name)
  58. os.environ["CUDA_VISIBLE_DEVICES"] = str(opt.GPU)
  59. batch_size = opt.batch_size
  60. lr = opt.lr
  62. name_list, noise_img_all, coordinate_list, stack_index = train_preprocess_lessMemoryMulStacks(opt)
  63. # print('name_list -----> ',name_list)
  64. ########################################################################################################################
  65. L1_pixelwise = torch.nn.L1Loss()
  66. L2_pixelwise = torch.nn.MSELoss()
  68. denoise_generator = Network_3D_Unet(in_channels = 1,
  69.                                     out_channels = 1,
  70.                                     f_maps=opt.fmap,
  71.                                     final_sigmoid = True)
  73. if torch.cuda.is_available():
  74.     denoise_generator = denoise_generator.cuda()
  75.     denoise_generator = nn.DataParallel(denoise_generator, device_ids=range(opt.ngpu))
  76.     print('\033[1;31mUsing {} GPU for training -----> \033[0m'.format(torch.cuda.device_count()))
  77.     L2_pixelwise.cuda()
  78.     L1_pixelwise.cuda()
  79. ########################################################################################################################
  80. optimizer_G = torch.optim.Adam(denoise_generator.parameters(),
  81.                                 lr=opt.lr, betas=(opt.b1, opt.b2))
  83. ########################################################################################################################
  84. cuda = True if torch.cuda.is_available() else False
  85. Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
  86. prev_time = time.time()
  88. ########################################################################################################################
  89. time_start=time.time()
  91. # start training
  92. for epoch in range(0, opt.n_epochs):
  93.     train_data = trainset(name_list, coordinate_list, noise_img_all,stack_index)
  94.     trainloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=4)
  95.     for iteration, (input, target) in enumerate(trainloader):
  96.         input=input.cuda()
  97.         target = target.cuda()
  98.         real_A=input
  99.         real_B=target
  100.         real_A = Variable(real_A)
  101.         #print('real_A shape -----> ', real_A.shape)
  102.         #print('real_B shape -----> ',real_B.shape)
  103.         fake_B = denoise_generator(real_A)
  104.         L1_loss = L1_pixelwise(fake_B, real_B)
  105.         L2_loss = L2_pixelwise(fake_B, real_B)
  106.         ################################################################################################################
  107.         optimizer_G.zero_grad()
  108.         # Total loss
  109.         Total_loss =  0.5*L1_loss + 0.5*L2_loss
  110.         Total_loss.backward()
  111.         optimizer_G.step()
  112.         ################################################################################################################
  113.         batches_done = epoch * len(trainloader) + iteration
  114.         batches_left = opt.n_epochs * len(trainloader) - batches_done
  115.         time_left = datetime.timedelta(seconds=int(batches_left * (time.time() - prev_time)))
  116.         prev_time = time.time()
  118.         if iteration%1 == 0:
  119.             time_end=time.time()
  120.             print('\r[Epoch %d/%d] [Batch %d/%d] [Total loss: %.2f, L1 Loss: %.2f, L2 Loss: %.2f] [ETA: %s] [Time cost: %.2d s]        ' 
  121.             % (
  122.                 epoch+1,
  123.                 opt.n_epochs,
  124.                 iteration+1,
  125.                 len(trainloader),
  126.                 Total_loss.item(),
  127.                 L1_loss.item(),
  128.                 L2_loss.item(),
  129.                 time_left,
  130.                 time_end-time_start
  131.             ), end=' ')
  133.         if (iteration+1)%len(trainloader) == 0:
  134.             print('\n', end=' ')
  136.         ################################################################################################################
  137.         # save model
  138.         if (iteration + 1) % (len(trainloader)) == 0:
  139.             model_save_name = pth_path + '//E_' + str(epoch+1).zfill(2) + '_Iter_' + str(iteration+1).zfill(4) + '.pth'
  140.             if isinstance(denoise_generator, nn.DataParallel): 
  141.                 torch.save(denoise_generator.module.state_dict(), model_save_name)  # parallel
  142.             else:
  143.                 torch.save(denoise_generator.state_dict(), model_save_name)         # not parallel