diff --git a/model_server/extensions/chaeo/params.py b/model_server/extensions/chaeo/params.py index 9146d435a5c6c3b41c33fee0f133ccb9cdc30e58..0880ca123df6e7eec9ce6235bbbba9ccaa1411cc 100644 --- a/model_server/extensions/chaeo/params.py +++ b/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] diff --git a/model_server/extensions/chaeo/zmask.py b/model_server/extensions/chaeo/zmask.py index 1111510c68019fdbe173b20ad6f0dfeae843f036..a457516b17cad744797af6305aa444f34d514e66 100644 --- a/model_server/extensions/chaeo/zmask.py +++ b/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(