added flips to transforms

JJBT · Sep 16, 2022 · d635d95 · d635d95
1 parent 792b89c
commit d635d95
Showing 1 changed file with 14 additions and 10 deletions.
diff --git a/utils/data_utils.py b/utils/data_utils.py
@@ -9,6 +9,10 @@ def get_msd_vanilla_transforms(
  roi_size: int,
  spacing: Sequence[float],
  modality: int or Sequence = 0,
+ RandFlipd_prob: float = 0.2,
+ RandRotate90d_prob: float = 0.2,
+ RandScaleIntensityd_prob: float = 0.1,
+ RandShiftIntensityd_prob: float = 0.1,
  **kwargs,
 ):
  modality = [modality] if isinstance(modality, int) else modality # modality = 0
@@ -25,16 +29,12 @@ def get_msd_vanilla_transforms(
  transforms.RandSpatialCropd(keys=["image", "label"],
  roi_size=(roi_size, roi_size, roi_size), random_size=False),
  transforms.Resized(keys=["image", "label"], spatial_size=(roi_size, roi_size, roi_size)),
- # transforms.RandCropByPosNegLabeld(
- # keys=["image", "label"],
- # label_key="label",
- # spatial_size=(roi_size, roi_size, roi_size),
- # pos=1,
- # neg=1,
- # num_samples=4,
- # image_key="image",
- # image_threshold=0,
- # ),
+ transforms.RandFlipd(keys=["image", "label"], prob=RandFlipd_prob, spatial_axis=0),
+ transforms.RandFlipd(keys=["image", "label"], prob=RandFlipd_prob, spatial_axis=1),
+ transforms.RandFlipd(keys=["image", "label"], prob=RandFlipd_prob, spatial_axis=2),
+ transforms.RandRotate90d(keys=["image", "label"], prob=RandRotate90d_prob, max_k=3),
+ transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=RandScaleIntensityd_prob),
+ transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=RandShiftIntensityd_prob),
  transforms.ToTensord(keys=["image", "label"]),
  ]
 
@@ -78,6 +78,10 @@ def get_loader(
  roi_size=roi_size,
  spacing=spacing,
  modality=modality,
+ RandFlipd_prob=RandFlipd_prob,
+ RandRotate90d_prob=RandRotate90d_prob,
+ RandScaleIntensityd_prob=RandScaleIntensityd_prob,
+ RandShiftIntensityd_prob=RandShiftIntensityd_prob,
  a_min=a_min,
  a_max=a_max,
  b_min=b_min,