|
- import sys
- sys.path.append('./')
- import torch
- import torchio as tio
- from torchio import Queue
- from torchio.transforms import (
- RandomFlip,
- RandomAffine,
- RandomElasticDeformation,
- RandomNoise,
- RandomMotion,
- RandomBiasField,
- RescaleIntensity,
- Resample,
- ToCanonical,
- ZNormalization,
- CropOrPad,
- HistogramStandardization,
- OneOf,
- Compose,
- )
- from pathlib import Path
-
- from hparam import hparams as hp
-
-
- class MedData_train(torch.utils.data.Dataset):
- def __init__(self, images_dir, labels_dir, args):
-
- self.aug = args.aug
- self.mode = args.mode
- self.patch_size = args.patch_size
- self.crop_or_pad_size = args.crop_or_pad_size
- self.fold_arch = args.fold_arch
-
- queue_length = 1088
- samples_per_volume = 20
-
- self.subjects = []
-
- images_dir = Path(images_dir)
- self.image_paths = sorted(images_dir.glob(self.fold_arch))
- labels_dir = Path(labels_dir)
- self.label_paths = sorted(labels_dir.glob(self.fold_arch))
-
- for (image_path, label_path) in zip(self.image_paths, self.label_paths):
- subject = tio.Subject(
- source=tio.ScalarImage(image_path),
- label=tio.LabelMap(label_path),
- )
- self.subjects.append(subject)
-
- self.transforms = self.transform()
-
- self.training_set = tio.SubjectsDataset(self.subjects, transform=self.transforms)
-
- probabilities = {0: 0.5, 1: 0.5}
- sampler = tio.data.LabelSampler(
- patch_size=self.patch_size,
- label_name='label',
- label_probabilities=probabilities,
- )
-
- self.queue_dataset = Queue(
- self.training_set,
- queue_length,
- samples_per_volume,
- sampler,
- num_workers=8
- )
-
-
- def transform(self):
-
- if self.mode == '3d':
- if self.aug:
- training_transform = Compose([
- ToCanonical(),
- Resample('source'),
- CropOrPad(self.crop_or_pad_size, padding_mode='reflect', mask_name='label'),
- # RandomMotion(),
- # RandomBiasField(),
- # ZNormalization(),
- # RandomNoise(),
- # RandomFlip(axes=(0,)),
- # OneOf({
- # RandomAffine(): 0.8,
- # RandomElasticDeformation(): 0.2,
- # }),
- ])
- else:
- training_transform = Compose([
- ToCanonical(),
- Resample('source'),
- CropOrPad(self.crop_or_pad_size, padding_mode='reflect', mask_name='label'),
- # ZNormalization(),
- ])
- elif self.mode == '2d':
- if self.aug:
- training_transform = Compose([
- ToCanonical(), # 对原图重排列使之转变为RAS方向轴排布(左->右,后->前,下->上)
- Resample('source'), # 重采样,改变图像像素的物理尺度
- CropOrPad(self.crop_or_pad_size, padding_mode='reflect', mask_name='label'), # 给定目标维度,如果原图大则裁剪,若小则填充
- # RandomMotion(), # 添加动态模糊,mri图像的噪声来源一部分是由于被测试者在采集时的动作导致的,通过当前增强来模拟该场景
- # RandomBiasField(), # 添加随机偏置场伪影,通常由于mri成像设备的磁场不均匀导致的低频强度变化(可以理解为一侧亮度偏低)
- # ZNormalization(), # 基于单张图像的归一化操作,即计算标准差和均值,将像素值强度的分布转化成高斯分布
- # RandomNoise(), # 添加随机高斯噪声
- # RandomFlip(axes=(0,)), # 给定中心轴(可以多个)翻转图像,可以使用数字(0,1,2)指定反转轴,也可以用字母(Left,Right,Height,Width)。允许设置翻转的概率,需要对每个反转轴单独指定。
- # OneOf({
- # RandomAffine(): 0.8, # 随机仿射变换,包括尺度(scale,需要指定缩放的比例,可以设置缩放时的差值策略),旋转(degrees,需要指定每个轴旋转的角度范围,可以设置旋转时pad的数值),平移(translation),还支持各向同性和设置以中心为基准进行变换。
- # RandomElasticDeformation(): 0.2, # 随机弹性形变,通过三次B样条插值实现位移,需要指定控制点的数量和最大位移距离,允许控制边缘是否形变
- # }) # OneOf:从给定增强变换序列中随机选择一项执行,允许设置每一项执行的概率
- ])
- else:
- training_transform = Compose([
- ToCanonical(),
- Resample('source'),
- CropOrPad(self.crop_or_pad_size, padding_mode='reflect', mask_name='label'),
- # ZNormalization(),
- ])
-
- else:
- raise Exception('no such kind of mode!')
-
-
- return training_transform
-
-
- class MedData_test(torch.utils.data.Dataset):
- def __init__(self, images_dir, labels_dir, args):
-
- queue_length = 1088
- samples_per_volume = 20
-
- self.crop_or_pad_size = args.crop_or_pad_size
- self.fold_arch = args.fold_arch
- self.patch_size = args.patch_size
- self.patch_overlap = args.patch_overlap
-
- self.subjects = []
-
- images_dir = Path(images_dir)
- self.image_paths = sorted(images_dir.glob(self.fold_arch))
- labels_dir = Path(labels_dir)
- self.label_paths = sorted(labels_dir.glob(self.fold_arch))
-
- for (image_path, label_path) in zip(self.image_paths, self.label_paths):
- subject = tio.Subject(
- source=tio.ScalarImage(image_path),
- label=tio.LabelMap(label_path),
- )
- self.subjects.append(subject)
-
- transform = Compose([
- ToCanonical(),
- Resample('source'),
- CropOrPad(self.crop_or_pad_size, padding_mode='reflect', mask_name='label'),
- # ZNormalization(),
- ])
-
- self.test_set = tio.SubjectsDataset(self.subjects, transform)
-
- probabilities = {0: 0.5, 1: 0.5}
- sampler = tio.data.LabelSampler(
- patch_size=self.patch_size,
- label_name='label',
- label_probabilities=probabilities,
- )
-
- self.queue_dataset = Queue(
- self.test_set,
- queue_length,
- samples_per_volume,
- sampler,
- num_workers=8
- )
|