Completed testing zmask module, including contour and filter optional args

extensions/chaeo/test_zstack.py

@@ -21,16 +21,17 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         mip = InMemoryDataAccessor(self.stack.get_one_channel_data(channel=0).data.max(axis=-1, keepdims=True))
         self.pxmap, result = pxmodel.infer(mip)
 
-        write_accessor_data_to_file(output_path / 'pxmap.tif', self.pxmap)
+        # write_accessor_data_to_file(output_path / 'pxmap.tif', self.pxmap)
         self.obmap = InMemoryDataAccessor(self.pxmap.data > pipeline_params['threshold'])
-        write_accessor_data_to_file(output_path / 'obmap.tif', self.obmap)
+        # write_accessor_data_to_file(output_path / 'obmap.tif', self.obmap)
 
-    def test_zmask_makes_correct_boxes(self):
+    def test_zmask_makes_correct_boxes(self, mask_type='boxes', filters=None):
         zmask, meta = build_stack_mask(
             'test_zmask_with boxes',
             self.obmap.get_one_channel_data(0),
             self.stack.get_one_channel_data(0),
-            mask_type='boxes',
+            mask_type=mask_type,
+            filters=filters,
         )
         zmask_acc = InMemoryDataAccessor(zmask)
         self.assertTrue(zmask_acc.is_mask())
@@ -51,4 +52,7 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         self.assertGreaterEqual(sh[0] * sh[1], ar)
 
     def test_zmask_makes_correct_contours(self):
-        pass
\ No newline at end of file
+        return self.test_zmask_makes_correct_boxes(mask_type='contours')
+
+    def test_zmask_makes_correct_boxes_with_filters(self):
+        return self.test_zmask_makes_correct_boxes(filters={'area': (1e3, 1e4)})
\ No newline at end of file

extensions/chaeo/zstack.py

@@ -98,8 +98,13 @@ def build_stack_mask(desc, obmap: GenericImageDataAccessor, stack: GenericImageD
     zi_st = np.zeros(stack.shape, dtype='bool')
     if mask_type == 'contours':
         zi_map = (lut[lamap] + 1.0).astype('int')
-        idxs = np.array([zi_map]) - 1
-        np.put_along_axis(zi_st, idxs, 1, axis=3)
+        idxs = np.array(zi_map) - 1
+        np.put_along_axis(
+            zi_st,
+            np.expand_dims(idxs, (2, 3)),
+            1,
+            axis=3
+        )
 
         # change background level from to 0 in final frame
         zi_st[:, :, :, -1][lamap == 0] = 0