Cleaning up zmask builder

extensions/chaeo/test_zstack.py

@@ -5,7 +5,7 @@ import numpy as np
 from conf.testing import output_path
 
 from extensions.chaeo.conf.testing import multichannel_zstack, pixel_classifier, pipeline_params
-from extensions.chaeo.zstack import build_stack_mask
+from extensions.chaeo.zmask import build_zmask_from_object_mask
 from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
 from extensions.ilastik.models import IlastikObjectClassifierModel, IlastikPixelClassifierModel
 
@@ -26,8 +26,7 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         # write_accessor_data_to_file(output_path / 'obmap.tif', self.obmap)
 
     def test_zmask_makes_correct_boxes(self, mask_type='boxes', filters=None):
-        zmask, meta = build_stack_mask(
-            'test_zmask_with boxes',
+        zmask, meta = build_zmask_from_object_mask(
             self.obmap.get_one_channel_data(0),
             self.stack.get_one_channel_data(0),
             mask_type=mask_type,

extensions/chaeo/zstack.py → extensions/chaeo/zmask.py

@@ -6,21 +6,27 @@ from skimage.measure import find_contours, label, regionprops_table
 from model_server.accessors import GenericImageDataAccessor
 
 # build a single boolean 3d mask (objects v. bboxes) and return bounding boxes
-def build_stack_mask(desc, obmap: GenericImageDataAccessor, stack: GenericImageDataAccessor, filters=None, mask_type='contour', expand_box_by=(0, 0)): # TODO: specify boxes data type
+def build_zmask_from_object_mask(
+        obmask: GenericImageDataAccessor,
+        stack: GenericImageDataAccessor,
+        filters=None,
+        mask_type='contour',
+        expand_box_by=(0, 0)
+):
     """
-
+    Given a 2D
     filters: dict of (min, max) tuples
     expand_box_by: (xy, z) pixelsf
     """
 
     # validate inputs
-    # assert len(stack.shape) == 3, stack.shape
     assert stack.chroma == 1
     assert stack.shape_dict['Z'] > 1
-    assert mask_type in ('contours', 'boxes'), mask_type # TODO: replace with call to validator
-
-    assert obmap.is_mask()
-    lamap = label(obmap.data[:, :, 0, 0])
+    assert mask_type in ('contours', 'boxes'), mask_type
+    assert obmask.is_mask()
+    assert obmask.chroma == 1
+    assert obmask.shape_dict['Z'] == 1
+    lamap = label(obmask.data[:, :, 0, 0])
 
     # build object query
     query_str = 'label > 0'  # always true