自己数据的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)