Hydra: Instantiating an object with a lambda function as an option

[oseymour]

This is a long one, sorry 😬.

The Goal

In my current setup, I use MONAI's Lambdad class as:

Lambdad(keys=["image", "mask"], func=lambda x: x.unsqueeze(-1) if x.ndim == 2 else x)

I'd like to be able to put my preprocessing steps in Hydra, so preproc.yaml would contain:

- _target_: monai.transforms.lambdad
  keys: ["image", "mask"]
  func: |
    lambda x: x.unsqueeze(-1) if x.ndim == 2 else x

as part of a list with all of my preprocessing steps. Full example preproc.yaml is below.

The Issue

Hydra instantiates the list without raising any errors, but when I try to use the instantiated transforms, I get the following errors (full traceback below).
The func arg is being initialized as a string, not callable.

TypeError: func must be None or callable but is str.

And Hydra is actually parsing the string incorrectly.

RuntimeError: applying transform <monai.transforms.utility.dictionary.Lambdad object at 0x000001C296533C80>
('lambda x: x.unsqueeze(-1) if x.ndim == 2 else x\n',
 'lambda x: x.unsqueeze(-1) if x.ndim == 2 else x\n')

Is what I'm trying to do possible with Hydra? Or can it not parse the YAML literal as Python and initialize the func arg with that code? Thanks in advance!

Full preproc.yaml file

common_pre_rand: &common_pre_rand
- _target_: monai.transforms.LoadImaged
  keys: ["image", "mask"]
  image_only: True
- _target_: monai.transforms.Lambdad
  keys: ["image", "mask"]
  func: | 
    lambda x: x.unsqueeze(-1) if x.ndim == 2 else x
- _target_: monai.transforms.EnsureChannelFirstd
  keys: ["image", "mask"]
  strict_check: True
  channel_dim: -1
- _target_: monai.transforms.Lambdad
  keys: "image"
  func: |
    lambda x: np.mean(x, axis=0, keepdims=True)
- _target_: monai.transforms.Rotate90d
  keys: ["image", "mask"]
  k: 3
- _target_: monai.transforms.Flipd
  keys: ["image", "mask"]
  spatial_axis: 1
- _target_: monai.transforms.ScaleIntensityd
  keys: "image"
- _target_: monai.transforms.AsDiscreted
  keys: "mask"
  threshold: 1
- _target_: monai.transforms.EnsureTyped
  keys: ["image", "mask"]


rand: &rand
- _target_: monai.transforms.RandScaleIntensityd
  keys: "image"
  prob: 0.5
  factors: 0.1
- _target_: monai.transforms.RandZoomd
  keys: ["image", "mask"]
  prob: 0.75
  min_zoom: 0.7
  max_zoom: 1.3
  mode: ["bilinear", "nearest"]
- _target_: monai.transforms.RandFlipd
  keys: ["image", "mask"]
  spatial_axis: 1
  prob: 0.5
- _target_: monai.transforms.RandAdjustContrastd
  keys: ["image", "mask"]
  prob: 1.0
  gamma: [0.75, 1.25]


common_post_rand: &common_post_rand
- _target_: monai.transforms.Resized
  keys: ["image", "mask"]
  mode: ["bilinear", "nearest"]
  spatial_size: [521, 512]


train_preproc:
- *common_pre_rand
- *rand
- *common_post_rand

val_preproc:
- *common_pre_rand
- *common_post_rand

Python code to instantiate and use preproc transforms with Hydra

from hydra import compose, initialize
from hydra.utils import instantiate
from monai.transforms import Compose

def parse_preproc():
    with initialize(version_base=None, config_path="."):
        cfg = compose(config_name="preproc")

    train_preproc = [step for group in cfg["train_preproc"] for step in group]
    train_preproc = instantiate(train_preproc)
    train_preproc = Compose(train_preproc)

    val_preproc = [step for group in cfg["val_preproc"] for step in group]
    val_preproc = instantiate(val_preproc)
    val_preproc = Compose(val_preproc)

    return train_preproc, val_preproc
    
train_preproc, val_preproc = parse_preproc()

transformed = train_preproc({
    "image": <<path to any jpg or png file>>,
    "mask": <<path to any jpg or png file>>
})

I can't provide example image and mask files since they're proprietary for work, sorry.

Full traceback

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File c:\Users\223113289.HCAD\AppData\Local\Programs\Python\Python312\Lib\site-packages\monai\transforms\transform.py:141, in apply_transform(transform, data, map_items, unpack_items, log_stats, lazy, overrides)
    [140](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:140)         return [_apply_transform(transform, item, unpack_items, lazy, overrides, log_stats) for item in data]
--> [141](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:141)     return _apply_transform(transform, data, unpack_items, lazy, overrides, log_stats)
    [142](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:142) except Exception as e:
    [143](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:143)     # if in debug mode, don't swallow exception so that the breakpoint
    [144](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:144)     # appears where the exception was raised.

File c:\Users\223113289.HCAD\AppData\Local\Programs\Python\Python312\Lib\site-packages\monai\transforms\transform.py:98, in _apply_transform(transform, data, unpack_parameters, lazy, overrides, logger_name)
     [96](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:96)     return transform(*data, lazy=lazy) if isinstance(transform, LazyTrait) else transform(*data)
---> [98](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:98) return transform(data, lazy=lazy) if isinstance(transform, LazyTrait) else transform(data)

File c:\Users\223113289.HCAD\AppData\Local\Programs\Python\Python312\Lib\site-packages\monai\transforms\utility\dictionary.py:1035, in Lambdad.__call__(self, data)
   [1034](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/utility/dictionary.py:1034) for key, func, overwrite in self.key_iterator(d, self.func, self.overwrite):
-> [1035](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/utility/dictionary.py:1035)     ret = self._lambd(img=d[key], func=func)
   [1036](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/utility/dictionary.py:1036)     if overwrite and isinstance(overwrite, bool):

File c:\Users\223113289.HCAD\AppData\Local\Programs\Python\Python312\Lib\site-packages\monai\transforms\utility\array.py:820, in Lambda.__call__(self, img, func)
    [819](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/utility/array.py:819) if not callable(fn):
--> [820](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/utility/array.py:820)     raise TypeError(f"func must be None or callable but is {type(fn).__name__}.")
    [821](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/utility/array.py:821) out = fn(img)

TypeError: func must be None or callable but is str.

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

RuntimeError                              Traceback (most recent call last)
Cell In[15], [line 1](vscode-notebook-cell:?execution_count=15&line=1)
----> [1](vscode-notebook-cell:?execution_count=15&line=1) train_preproc({
      [2](vscode-notebook-cell:?execution_count=15&line=2)     "image": "../data/carotid_channel_data_205/validation/images/case_005.dcm", 
      [3](vscode-notebook-cell:?execution_count=15&line=3)     "mask": "../data/carotid_channel_data_205/validation/masks/case_005.nii"
      [4](vscode-notebook-cell:?execution_count=15&line=4) })

File c:\Users\223113289.HCAD\AppData\Local\Programs\Python\Python312\Lib\site-packages\monai\transforms\compose.py:335, in Compose.__call__(self, input_, start, end, threading, lazy)
    [333](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:333) def __call__(self, input_, start=0, end=None, threading=False, lazy: bool | None = None):
    [334](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:334)     _lazy = self._lazy if lazy is None else lazy
--> [335](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:335)     result = execute_compose(
    [336](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:336)         input_,
    [337](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:337)         transforms=self.transforms,
    [338](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:338)         start=start,
    [339](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:339)         end=end,
    [340](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:340)         map_items=self.map_items,
    [341](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:341)         unpack_items=self.unpack_items,
    [342](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:342)         lazy=_lazy,
    [343](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:343)         overrides=self.overrides,
    [344](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:344)         threading=threading,
    [345](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:345)         log_stats=self.log_stats,
    [346](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:346)     )
    [348](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:348)     return result

File c:\Users\223113289.HCAD\AppData\Local\Programs\Python\Python312\Lib\site-packages\monai\transforms\compose.py:111, in execute_compose(data, transforms, map_items, unpack_items, start, end, lazy, overrides, threading, log_stats)
    [109](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:109)     if threading:
    [110](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:110)         _transform = deepcopy(_transform) if isinstance(_transform, ThreadUnsafe) else _transform
--> [111](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:111)     data = apply_transform(
    [112](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:112)         _transform, data, map_items, unpack_items, lazy=lazy, overrides=overrides, log_stats=log_stats
    [113](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:113)     )
    [114](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:114) data = apply_pending_transforms(data, None, overrides, logger_name=log_stats)
    [115](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/compose.py:115) return data

File c:\Users\223113289.HCAD\AppData\Local\Programs\Python\Python312\Lib\site-packages\monai\transforms\transform.py:171, in apply_transform(transform, data, map_items, unpack_items, log_stats, lazy, overrides)
    [169](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:169)     else:
    [170](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:170)         _log_stats(data=data)
--> [171](file:///C:/Users/223113289.HCAD/AppData/Local/Programs/Python/Python312/Lib/site-packages/monai/transforms/transform.py:171) raise RuntimeError(f"applying transform {transform}") from e

RuntimeError: applying transform <monai.transforms.utility.dictionary.Lambdad object at 0x000001C296533C80>
('lambda x: x.unsqueeze(-1) if x.ndim == 2 else x\n',
 'lambda x: x.unsqueeze(-1) if x.ndim == 2 else x\n')