Add missing Attachments WIP

Author: constantinpape

bioimageio/core/resource_io/nodes.py

@@ -139,36 +139,41 @@ def __call__(self, *args, **kwargs):
 
 
 @dataclass
-class KerasHdf5WeightsEntry(model_raw_nodes.KerasHdf5WeightsEntry):
+class KerasHdf5WeightsEntry(Node, model_raw_nodes.KerasHdf5WeightsEntry):
     source: Path = missing
 
 
 @dataclass
-class OnnxWeightsEntry(model_raw_nodes.OnnxWeightsEntry):
+class OnnxWeightsEntry(Node, model_raw_nodes.OnnxWeightsEntry):
     source: Path = missing
 
 
 @dataclass
-class PytorchStateDictWeightsEntry(model_raw_nodes.PytorchStateDictWeightsEntry):
+class PytorchStateDictWeightsEntry(Node, model_raw_nodes.PytorchStateDictWeightsEntry):
     source: Path = missing
     architecture: Union[_Missing, ImportedSource] = missing
 
 
 @dataclass
-class PytorchScriptWeightsEntry(model_raw_nodes.PytorchScriptWeightsEntry):
+class PytorchScriptWeightsEntry(Node, model_raw_nodes.PytorchScriptWeightsEntry):
     source: Path = missing
 
 
 @dataclass
-class TensorflowJsWeightsEntry(model_raw_nodes.TensorflowJsWeightsEntry):
+class TensorflowJsWeightsEntry(Node, model_raw_nodes.TensorflowJsWeightsEntry):
     source: Path = missing
 
 
 @dataclass
-class TensorflowSavedModelBundleWeightsEntry(model_raw_nodes.TensorflowSavedModelBundleWeightsEntry):
+class TensorflowSavedModelBundleWeightsEntry(Node, model_raw_nodes.TensorflowSavedModelBundleWeightsEntry):
     source: Path = missing
 
 
+@dataclass
+class Attachments(Node, model_raw_nodes.Attachments):
+    files: List[Path] = missing
+
+
 WeightsEntry = Union[
     KerasHdf5WeightsEntry,
     OnnxWeightsEntry,

tests/build_spec/test_build_spec.py

@@ -90,8 +90,8 @@ def _test_build_spec(
         loaded_config = loaded_model.config
         assert "deepimagej" in loaded_config
 
-    attachments = loaded_model.attachments or {}
-    if "files" in attachments:
+    attachments = loaded_model.attachments
+    if attachments is not missing and attachments.files is not missing:
         for attached_file in attachments["files"]:
             assert attached_file.exists()
 
@@ -132,3 +132,7 @@ def test_build_spec_tfjs(any_tensorflow_js_model, tmp_path):
 
 def test_build_spec_deepimagej(unet2d_nuclei_broad_model, tmp_path):
     _test_build_spec(unet2d_nuclei_broad_model, tmp_path / "model.zip", "pytorch_script", add_deepimagej_config=True)
+
+
+# def test_build_spec_deepimagej_keras(unet2d_keras, tmp_path):
+#     _test_build_spec(unet2d_keras, tmp_path / "model.zip", "pytorch_script", add_deepimagej_config=True)