Moved zmask generation into RoiSet method

model_server/extensions/chaeo/params.py

@@ -27,7 +27,6 @@ class RoiFilter(BaseModel):
 
 
 class RoiSetMetaParams(BaseModel):
-    mask_type: str = 'boxes'
     filters: Union[RoiFilter, None] = None
     expand_box_by: List[int] = [128, 0]
 

model_server/extensions/chaeo/zmask.py

@@ -32,12 +32,12 @@ class RoiSet(object):
             acc_raw: GenericImageDataAccessor,
             params: RoiSetMetaParams = RoiSetMetaParams(),
     ):
-        self.zmask, self.zmask_meta, self.df, self.interm = build_zmask_from_object_mask(
+        self.zmask_meta, self.df, self.interm = build_zmask_from_object_mask(
             acc_obj_ids,
             acc_raw,
             params=params,
         )
-
+        self.acc_obj_ids = acc_obj_ids
         self.acc_raw = acc_raw
         self.count = len(self.zmask_meta)
         self.object_id_labels = self.interm['label_map']
@@ -76,8 +76,41 @@ class RoiSet(object):
     def get_slices(self):
         return [zm.slice for zm in self.zmask_meta]
 
-    def get_zmask(self): # TODO: on-the-fly generation of zmask array
-        return self.zmask
+    def get_zmask(self, mask_type='boxes'):
+        """
+        Return a mask of same dimensionality as raw data
+
+        :param kwargs: variable-length keyword arguments
+            mask_type: if 'boxes', zmask is True in each object's complete bounding box; otherwise 'contours'
+        """
+
+        assert mask_type in ('contours', 'boxes')
+        zi_st = np.zeros(self.acc_raw.shape, dtype='bool')
+        lamap = self.acc_obj_ids
+
+        # make an object map where label is replaced by focus position in stack and background is -1
+        lut = np.zeros(lamap.max() + 1) - 1
+        lut[self.df.label] = self.df.zi
+
+        if mask_type == 'contours':
+            zi_map = (lut[lamap] + 1.0).astype('int')
+            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
+
+        elif mask_type == 'boxes':
+            for bb in self.zmask_meta:
+                sl = bb['slice']
+                zi_st[sl] = 1
+
+        return zi_st
 
     def classify_by(self, channel, object_classification_model: InstanceSegmentationModel):
         # do this on a patch basis, i.e. only one object per frame
@@ -170,10 +203,8 @@ def build_zmask_from_object_mask(
             argmax: np.ndarray (h x w x 1 x 1) z-index of highest intensity in zstack
     """
     filters = params.filters
-    mask_type = params.mask_type
     expand_box_by = params.expand_box_by
     # validate inputs
-    assert mask_type in ('contours', 'boxes'), mask_type
     assert zstack.hw == obmask.shape
 
     lamap = obmask
@@ -209,10 +240,6 @@ def build_zmask_from_object_mask(
     df['keeper'] = False
     df.loc[df.query(query_str).index, 'keeper'] = True
 
-    # make an object map where label is replaced by focus position in stack and background is -1
-    lut = np.zeros(lamap.max() + 1) - 1
-    lut[df.label] = df.zi
-
     # convert bounding boxes to numpy slice objects
     ebxy, ebz = expand_box_by
     h, w, c, nz = zstack.shape
@@ -253,33 +280,13 @@ def build_zmask_from_object_mask(
             'mask': mask
         })
 
-    # build mask z-stack # TODO: on-the-fly
-    zi_st = np.zeros(zstack.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,
-            np.expand_dims(idxs, (2, 3)),
-            1,
-            axis=3
-        )
-
-        # change background level from to 0 in final frame
-        zi_st[:, :, :, -1][lamap == 0] = 0
-
-    elif mask_type == 'boxes':
-        for bb in meta:
-            sl = bb['slice']
-            zi_st[sl] = 1
-
     # return intermediate image arrays
     interm = {
         'label_map': lamap,
         'argmax': argmax,
     }
 
-    return zi_st, meta, df, interm
+    return meta, df, interm
 
 
 def project_stack_from_focal_points(