Skip to content
Snippets Groups Projects
Commit ed05746d authored by s193396's avatar s193396
Browse files

fixing augmentation errors

parent 24923e92
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
qim3d/ml/_augmentations.py
+ 3
− 3
View file @ ed05746d
...@@ -86,11 +86,11 @@ class Augmentation: ...@@ -86,11 +86,11 @@ class Augmentation:
level_aug = [] level_aug = []
elif level == 'light': elif level == 'light':
level_aug = [RandRotate90(prob=1, spatial_axes=(0, 1, 2))] if self.is_3d else [RandRotate90(prob=1)] level_aug = [RandRotate90(prob=1, spatial_axes=(0, 1))] if self.is_3d else [RandRotate90(prob=1)]
elif level == 'moderate': elif level == 'moderate':
level_aug = [ level_aug = [
RandRotate90(prob=1, spatial_axes=(0, 1, 2)) if self.is_3d else RandRotate90(prob=1), RandRotate90(prob=1, spatial_axes=(0, 1)) if self.is_3d else RandRotate90(prob=1),
RandFlip(prob=0.3, spatial_axis=0), RandFlip(prob=0.3, spatial_axis=0),
RandFlip(prob=0.3, spatial_axis=1), RandFlip(prob=0.3, spatial_axis=1),
RandGaussianSmooth(sigma_x=(0.7, 0.7), prob=0.1), RandGaussianSmooth(sigma_x=(0.7, 0.7), prob=0.1),
...@@ -99,7 +99,7 @@ class Augmentation: ...@@ -99,7 +99,7 @@ class Augmentation:
elif level == 'heavy': elif level == 'heavy':
level_aug = [ level_aug = [
RandRotate90(prob=1, spatial_axes=(0, 1, 2)) if self.is_3d else RandRotate90(prob=1), RandRotate90(prob=1, spatial_axes=(0, 1)) if self.is_3d else RandRotate90(prob=1),
RandFlip(prob=0.7, spatial_axis=0), RandFlip(prob=0.7, spatial_axis=0),
RandFlip(prob=0.7, spatial_axis=1), RandFlip(prob=0.7, spatial_axis=1),
RandGaussianSmooth(sigma_x=(1.2, 1.2), prob=0.3), RandGaussianSmooth(sigma_x=(1.2, 1.2), prob=0.3),
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment