How to Add Custom Policies and New Augmentation Functions

I'd like to add a new operation to the list of available functions to a policy. In this case, the random gamma transform in the albumentations library. As far as I know to do this you need to define the operation in policy_operations.py

class RandomGamma(Operation):
    def __init__(self, temperature):
        super().__init__(temperature, requires_uint8_scaling=True)

    def apply_operation(self, input, value):
        return F.gamma_transform(input, gamma_limit=value)

    def as_transform(self, value, p):
        return albumentations.RandomGamma(p=p, gamma_limit=(1, value))

And then define it again in functional.py in albumentations_pytorch

def gamma_transform(img_batch, gamma):
    return torch.pow(img_batch, gamma)

You then need to create a new policy yaml file which has the operation in it's list:

operations:
  - _target_: autoalbument.faster_autoaugment.models.policy_operations.RandomGamma

This policy model is then referenced by the config file in ~\cli\conf\config.yaml

Having done all that, I get an error. Please find the full traceback below:

Traceback (most recent call last):
  File "/home/eas90/anaconda3/envs/AutoAug2/lib/python3.9/site-packages/hydra/_internal/utils.py", line 544, in _locate
    import_module(mod)
  File "/home/eas90/anaconda3/envs/AutoAug2/lib/python3.9/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1030, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1007, in _find_and_load
  File "<frozen importlib._bootstrap>", line 981, in _find_and_load_unlocked
ModuleNotFoundError: No module named 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma'; 'autoalbument.faster_autoaugment.models.policy_operations' is not a package

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/eas90/anaconda3/envs/AutoAug2/lib/python3.9/site-packages/hydra/utils.py", line 61, in call
    type_or_callable = _locate(cls)
  File "/home/eas90/anaconda3/envs/AutoAug2/lib/python3.9/site-packages/hydra/_internal/utils.py", line 546, in _locate
    raise ImportError(
ImportError: Encountered error: `No module named 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma'; 'autoalbument.faster_autoaugment.models.policy_operations' is not a package` when loading module 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/eas90/anaconda3/envs/AutoAug2/lib/python3.9/site-packages/hydra/utils.py", line 63, in call
    return _instantiate_class(type_or_callable, config, *args, **kwargs)
  File "/home/eas90/anaconda3/envs/AutoAug2/lib/python3.9/site-packages/hydra/_internal/utils.py", line 500, in _instantiate_class
    return clazz(*args, **final_kwargs)
  File "/zpool-00/home/eas90/PycharmProjects/AutoAug/autoalbument/autoalbument/faster_autoaugment/search.py", line 17, in __init__
    self.model = self.create_model()
  File "/zpool-00/home/eas90/PycharmProjects/AutoAug/autoalbument/autoalbument/faster_autoaugment/search.py", line 33, in create_model
    return FAAClassificationModel(cfg)
  File "/zpool-00/home/eas90/PycharmProjects/AutoAug/autoalbument/autoalbument/faster_autoaugment/models/faa_model.py", line 22, in __init__
    self.policy_model = self.create_policy_model()
  File "/zpool-00/home/eas90/PycharmProjects/AutoAug/autoalbument/autoalbument/faster_autoaugment/models/faa_model.py", line 38, in create_policy_model
    policy_operations = [
  File "/zpool-00/home/eas90/PycharmProjects/AutoAug/autoalbument/autoalbument/faster_autoaugment/models/faa_model.py", line 39, in <listcomp>
    instantiate(operation, temperature=policy_model_cfg.temperature)
  File "/home/eas90/anaconda3/envs/AutoAug2/lib/python3.9/site-packages/hydra/utils.py", line 70, in call
    raise HydraException(f"Error calling '{cls}' : {e}") from e
hydra.errors.HydraException: Error calling 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma' : Encountered error: `No module named 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma'; 'autoalbument.faster_autoaugment.models.policy_operations' is not a package` when loading module 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/zpool-00/home/eas90/PycharmProjects/AutoAug/autoalbument/autoalbument/cli/search.py", line 54, in main
    searcher = instantiate(cfg.searcher, cfg=cfg)
  File "/home/eas90/anaconda3/envs/AutoAug2/lib/python3.9/site-packages/hydra/utils.py", line 70, in call
    raise HydraException(f"Error calling '{cls}' : {e}") from e
hydra.errors.HydraException: Error calling 'autoalbument.faster_autoaugment.search.FasterAutoAugmentSearcher' : Error calling 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma' : Encountered error: `No module named 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma'; 'autoalbument.faster_autoaugment.models.policy_operations' is not a package` when loading module 'autoalbument.faster_autoaugment.models.policy_operations.RandomGamma'

Has anyone tried to add operations into the function pool before and know what to do? I'm very stuck. Thank you for reading.

albumentations-team / autoalbument

How to Add Custom Policies and New Augmentation Functions #27