TotalSegmentator fails with "ImportError: cannot import name 'ResidualEncoderUNet''

[lassoan]

Recent update of dynamic_network_architectures package contains a backward-incompatible change. ResidualEncoderUNet class is in a different module now:

In dynamic_network_architectures version 0.2:

from dynamic_network_architectures.architectures.unet import ResidualEncoderUNet

In dynamic_network_architectures version 0.4:

from dynamic_network_architectures.architectures.residual_unet import ResidualEncoderUNet

This change breaks nnunet, which breaks totalsegmentator.

Could you update totalsegmentator requirements specification to avoid this incompatible dynamic_network_architectures version until the issue is properly sorted out? For a proper solution, maintainers of nnunet and dynamic_network_architectures need to talk.

Full stack trace:

Traceback (most recent call last):
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Scripts\TotalSegmentator.exe\__main__.py", line 7, in <module>
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\totalsegmentator\bin\TotalSegmentator.py", line 127, in main
    totalsegmentator(args.input, args.output, args.ml, args.nr_thr_resamp, args.nr_thr_saving,
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\totalsegmentator\python_api.py", line 293, in totalsegmentator
    seg_img, ct_img = nnUNet_predict_image(input, output, task_id, model=model, folds=folds,
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\totalsegmentator\nnunet.py", line 375, in nnUNet_predict_image
    nnUNetv2_predict(tmp_dir, tmp_dir, tid, model, folds, trainer, tta,
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\totalsegmentator\nnunet.py", line 178, in nnUNetv2_predict
    predict_from_raw_data(dir_in,
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\nnunetv2\inference\predict_from_raw_data.py", line 159, in predict_from_raw_data
    load_what_we_need(model_training_output_dir, use_folds, checkpoint_name)
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\nnunetv2\inference\predict_from_raw_data.py", line 96, in load_what_we_need
    trainer_class = recursive_find_python_class(join(nnunetv2.__path__[0], "training", "nnUNetTrainer"),
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\nnunetv2\utilities\find_class_by_name.py", line 12, in recursive_find_python_class
    m = importlib.import_module(current_module + "." + modname)
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\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 986, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 680, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 850, in exec_module
  File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\nnunetv2\training\nnUNetTrainer\nnUNetTrainer.py", line 51, in <module>
    from nnunetv2.utilities.get_network_from_plans import get_network_from_plans
  File "C:\Users\Logan\AppData\Local\slicer.org\Slicer 5.6.1\lib\Python\Lib\site-packages\nnunetv2\utilities\get_network_from_plans.py", line 1, in <module>
    from dynamic_network_architectures.architectures.unet import PlainConvUNet, ResidualEncoderUNet
ImportError: cannot import name 'ResidualEncoderUNet' from 'dynamic_network_architectures.architectures.unet' 

[jamesobutler]

It appears that nnunetv2 version 2.3 specifically put an upper pin on dynamic_network_architectures due to this issue. https://github.com/MIC-DKFZ/nnUNet/commit/9fd926b7ba161085117ba2e55ce3b1f84fa68245. Therefore TotalSegmentator should probably enforce nnunetv2 >=2.3 instead of >=2.2.1. https://github.com/wasserth/TotalSegmentator/blob/26318286d9b184e9b8aed340d5faeda97b705725/setup.py#L24

@lassoan SlicerTotalSegmentator specifies to use TotalSegmentator v2.0.5 which specifies nnunetv2==2.1 which has a dynamic-network-architectures>=0.2 specification. Upgrading to TotalSegmentator 2.1.0 would result in it using nnunetv2>=2.2.1 which would grab nnunetv2 version v2.3 if a fresh install of TotalSegmentator and would pull dynamic-network-architectures<=0.3 to avoid issues. Unfortunately it seems that nnUNetv2 probably still needed their previous pin of dynamic-network-architectures>=0.2 to enforce a lower bound as well. nnUNetv2 might not fix that as it appears their latest master branch probably has some compatibility fixes as it enforces dynamic-network-architectures>=0.4. It's unclear when they will release a new version.

[lassoan]

The issue was just fixed in nnunet yesterday: https://github.com/MIC-DKFZ/nnUNet/commit/f569e34d0265723288a64eca579609e0274b1a0b - hopefully a release comes soon, then TotalSegmentator can be updated to use the new nnunet, and then Slicer can be updated to use the new TotalSegmentator.

[lassoan]

I've submitted an issue to dynamic-network-architectures to let the maintainers know that this is a major issue and asked for yanking the incompatible dynamic-network-architectures versions. We'll see.

[pwaver]

Apparent same problem today trying to install U-Mamba:

from dynamic_network_architectures.architectures.unet import PlainConvUNet, ResidualEncoderUNet ImportError: cannot import name 'ResidualEncoderUNet' from 'dynamic_network_architectures.architectures.unet' (/opt/conda/envs/umamba/lib/python3.10/site-packages/dynamic_network_architectures/architectures/unet.py)

Bill