Test that calls RoiSet directly

model_server/extensions/chaeo/tests/test_zstack.py

@@ -8,7 +8,7 @@ from model_server.extensions.chaeo.conf.testing import multichannel_zstack, pixe
 from extensions.chaeo.params import RoiSetExportParams, RoiSetMetaParams
 from model_server.extensions.chaeo.products import export_patches_from_zstack, export_multichannel_patches_from_zstack, export_patch_masks_from_zstack
 from model_server.extensions.chaeo.workflows import infer_object_map_from_zstack
-from model_server.extensions.chaeo.zmask import build_zmask_from_object_mask, get_label_ids
+from model_server.extensions.chaeo.zmask import build_zmask_from_object_mask, get_label_ids, RoiSet
 from model_server.base.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
 from model_server.extensions.ilastik.models import IlastikPixelClassifierModel
 from model_server.base.models import DummyInstanceSegmentationModel
@@ -42,7 +42,7 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
 
     def test_zmask_makes_correct_boxes(self, mask_type='boxes', **kwargs):
         id_map = get_label_ids(self.seg_mask)
-        zmask, meta, df, interm = build_zmask_from_object_mask(
+        zmask, meta, _, interm = build_zmask_from_object_mask(
             id_map,
             self.stack_ch_pa,
             params=RoiSetMetaParams(mask_type=mask_type, filters=kwargs.get('filters')),
@@ -75,12 +75,10 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         acc_zstack_slice = InMemoryDataAccessor(self.stack_ch_pa.data[:, :, :, 0])
         self.assertEqual(acc_zstack_slice.nz, 1)
         id_map = get_label_ids(self.seg_mask)
-        zmask, meta, df, interm = build_zmask_from_object_mask(
-            id_map,
-            acc_zstack_slice,
-            params=RoiSetMetaParams(mask_type='boxes'),
-            **kwargs,
-        )
+
+        roiset = RoiSet(id_map, acc_zstack_slice, params=RoiSetMetaParams(mask_type='boxes'))
+        zmask = roiset.get_zmask()
+
         zmask_acc = InMemoryDataAccessor(zmask)
         self.assertTrue(zmask_acc.is_mask())
 
@@ -120,7 +118,7 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
 
     def test_flatten_image(self):
         id_map = get_label_ids(self.seg_mask)
-        zmask, meta, df, interm = build_zmask_from_object_mask(
+        _, _, df, _ = build_zmask_from_object_mask(
             id_map,
             self.stack_ch_pa,
             params=RoiSetMetaParams(mask_type='boxes')

model_server/extensions/chaeo/zmask.py

@@ -86,6 +86,9 @@ class RoiSet(object):
     def get_object_mask_by_class(self, class_id):
         return self.object_id_labels == class_id
 
+    def get_zmask(self): # TODO: on-the-fly generation of zmask array
+        return self.zmask
+
     def run_exports(self, where, channel, prefix, params: RoiSetExportParams):
         if not self.count:
             return
@@ -161,17 +164,9 @@ def build_zmask_from_object_mask(
     mask_type = params.mask_type
     expand_box_by = params.expand_box_by
     # validate inputs
-    # assert zstack.chroma == 1
     assert mask_type in ('contours', 'boxes'), mask_type
-    # assert obmask.is_mask()
-    # assert obmask.chroma == 1
-    # assert obmask.nz == 1
-    if zstack.hw != obmask.shape:
-        input()
     assert zstack.hw == obmask.shape
 
-    # assign object labels and build object query
-    # lamap = label(obmask.data[:, :, 0, 0]).astype('uint16')
     lamap = obmask
     query_str = 'label > 0'  # always true
     if filters is not None: