Coverage for bioimageio/core/weight_converters/_utils_onnx.py: 100%

1from collections import defaultdict 

2from itertools import chain 

3from typing import DefaultDict, Dict 

4 

5from bioimageio.spec.model.v0_5 import ModelDescr 

6 

7 

8def get_dynamic_axes(model_descr: ModelDescr): 

9 dynamic_axes: DefaultDict[str, Dict[int, str]] = defaultdict(dict) 

10 for d in chain(model_descr.inputs, model_descr.outputs): 

11 for i, ax in enumerate(d.axes): 

12 if not isinstance(ax.size, int): 

13 dynamic_axes[str(d.id)][i] = str(ax.id) 

14 

15 return dynamic_axes