feat(Zettaset): enable default volume for user-specified annotations

Author: torms3

deepem/data/dataset/multi_zettaset.py

@@ -130,6 +130,7 @@ def load_sample(
     padding: tuple[int, int, int] = (0, 0, 0),
     no_mask: bool = False,
     zettaset_lookup: dict[str, str] | None = None,
+    zettaset_default: dict[str, str] = {},
     zettaset_resolution: tuple[int, int, int] | None = None,
     requires_binarize: list[str] = [],
     **kwargs
@@ -168,8 +169,17 @@ def convert_array(arr: ArrayLike) -> np.ndarray:
     for name, key in zettaset_lookup.items():
 
         # Annotation
-        vol = sample.read(key)[key]
-        dset[name] = convert_array(vol)
+        if key in sample.annotation_names:
+            vol = sample.read(key)[key]
+            dset[name] = convert_array(vol)
+        else:
+            if key in zettaset_default:
+                shape = tuple(map(int, bbox.size()))
+                dtype = zettaset_default[key]
+                vol = np.zeros(shape, dtype=dtype)
+                dset[name] = np.transpose(vol, (2, 1, 0))
+            else:
+                raise KeyError(f"Annotation '{key}' not found.")
         anno_log = f"\t{name}: {dset[name].shape}"
 
         # Binarize

deepem/train/option.py

@@ -28,6 +28,7 @@ def initialize(self):
         self.parser.add_argument('--zettaset_path', type=str, default=[], nargs='+')
         self.parser.add_argument('--zettaset_specs', type=json.loads, default={})
         self.parser.add_argument('--zettaset_lookup', type=json.loads, default=None)
+        self.parser.add_argument('--zettaset_default', type=json.loads, default={})
         self.parser.add_argument('--zettaset_padding', type=vec3, default=(0, 0, 0))
         self.parser.add_argument('--zettaset_padding_spec', type=json.loads, default={})
         self.parser.add_argument('--zettaset_resolution', type=vec3f, default=None)
@@ -314,6 +315,7 @@ def parse(self):
             class_keys=class_keys,
             glia_mask=opt.glia_mask,
             zettaset_lookup=opt.zettaset_lookup,
+            zettaset_default=opt.zettaset_default,
             zettaset_padding=opt.zettaset_padding,
             zettaset_padding_spec=opt.zettaset_padding_spec,
             zettaset_resolution=opt.zettaset_resolution,