diff --git a/dacapo/experiments/trainers/gp_augments/__init__.py b/dacapo/experiments/trainers/gp_augments/__init__.py index 0c93d4603..c5698a851 100644 --- a/dacapo/experiments/trainers/gp_augments/__init__.py +++ b/dacapo/experiments/trainers/gp_augments/__init__.py @@ -4,3 +4,4 @@ from .gamma_config import GammaAugmentConfig from .intensity_config import IntensityAugmentConfig from .intensity_scale_shift_config import IntensityScaleShiftAugmentConfig +from .noise_config import NoiseAugmentConfig \ No newline at end of file diff --git a/dacapo/experiments/trainers/gp_augments/noise_config.py b/dacapo/experiments/trainers/gp_augments/noise_config.py new file mode 100644 index 000000000..4e9dc1000 --- /dev/null +++ b/dacapo/experiments/trainers/gp_augments/noise_config.py @@ -0,0 +1,34 @@ +from .augment_config import AugmentConfig + +import gunpowder as gp + +import attr + +from typing import Tuple + + +@attr.s +class NoiseAugmentConfig(AugmentConfig): + """ + This class manages the configuration of gamma augmentation for a given dataset. + + Attributes: + gamma_range: A tuple of float values represents the min and max range of gamma noise + to apply on the raw data. + Methods: + node(): Constructs a node in the augmentation pipeline. + """ + + kwargs: dict[str, any] = attr.ib( + metadata={ + "help_text": "key word arguments for skimage `random_noise`. For more details see " + "https://scikit-image.org/docs/stable/api/skimage.util.html#skimage.util.random_noise" + }, + factory=lambda: dict(), + ) + + def node(self, raw_key: gp.ArrayKey, _gt_key=None, _mask_key=None): + """ + Constructs a node in the augmentation pipeline. + """ + return gp.NoiseAugment(raw_key, **self.kwargs)