Pytorch dataset类

自己数据的dataset类，传入数据

比如训练数据格式为.txt格式，字符串类型如下所示：

/home/xx/data/volume_01.nii.gz /home/xx/data/label_01.nii.gz 
/home/xx/data/volume_02.nii.gz /home/xx/data/label_02.nii.gz 
/home/xx/data/volume_03.nii.gz /home/xx/data/label_03.nii.gz 
/home/xx/data/volume_04.nii.gz /home/xx/data/label_04.nii.gz 
/home/xx/data/volume_05.nii.gz /home/xx/data/label_05.nii.gz 
...

• Dataset类构造

from torch.utils.data import Dataset as dataset
import os
import SimpleITK as sitk
import numpy as np
import torch


class Train_Dataset(dataset):
    def __init__(self, args):
    	self.args = args
    	self.filename_list = self.load_file_name_list(os.path.join(args.dataset_path, 'train_path_list.txt'))
    	# self.transforms 也可以在__init__中加入一些数据增强方法

	def __getitem__(self, index):
		ct = sitk.ReadImage(self.filename_list[index][0], sitk.sitkInt16)# medical method
        seg = sitk.ReadImage(self.filename_list[index][1], sitk.sitkUInt8)
		
		ct_array = sitk.GetArrayFromImage(ct)
        seg_array = sitk.GetArrayFromImage(seg)
        
		ct_array = np.clip((ct_array - 500) / self.args.norm_factor, -1, 1) # clip and normalization method
		ct_array = ct_array.astype(np.float32)
		ct_array = torch.FloatTensor(ct_array).unsqueeze(0)# to tensor ?
        seg_array = torch.FloatTensor(seg_array).unsqueeze(0)
		return ct_array, seg_array.squeeze(0)
		
	def __len__(self):
		return len(self.filename_list) # return dataset size

• Dataloader载入数据

from dataset.dataset_lits_train import Train_Dataset
from torch.utils.data import DataLoader

train_loader = DataLoader(dataset=Train_Dataset(args),batch_size=args.batch_size,num_workers=args.n_threads, shuffle=True)