diff --git a/model_server/extensions/chaeo/tests/test_zstack.py b/model_server/extensions/chaeo/tests/test_zstack.py index 70daa788a41ae3d2b31667012fb0512a6cfd681b..cc9a4297df91e92ad2985e88bd7ad4f1984d9857 100644 --- a/model_server/extensions/chaeo/tests/test_zstack.py +++ b/model_server/extensions/chaeo/tests/test_zstack.py @@ -94,6 +94,26 @@ class TestZStackDerivedDataProducts(unittest.TestCase): def test_zmask_makes_correct_expanded_boxes(self): return self.test_zmask_makes_correct_boxes(expand_box_by=(64, 2)) + def test_zmask_slices_are_valid(self): + roiset = self.test_zmask_makes_correct_boxes() + slices = [roiset.get_slice_at(i) for i in roiset.get_df().index] + for s in slices: + ebb = roiset.acc_raw.data[s] + self.assertEqual(len(ebb.shape), 4) + self.assertTrue(np.all([si >= 1 for si in ebb.shape])) + + def test_zmask_rel_slices_are_valid(self): + roiset = self.test_zmask_makes_correct_boxes() + slices = [roiset.get_slice_at(i) for i in roiset.get_df().index] + rel_slices = [roiset.get_rel_slice_at(i) for i in roiset.get_df().index] + for i, s in enumerate(slices): + ebb = roiset.acc_raw.data[s] + self.assertEqual(len(ebb.shape), 4) + self.assertTrue(np.all([si >= 1 for si in ebb.shape])) + rbb = ebb[rel_slices[i]] + self.assertEqual(len(rbb.shape), 4) + self.assertTrue(np.all([si >= 1 for si in rbb.shape])) + def test_make_2d_patches_from_zmask(self): roiset = self.test_zmask_makes_correct_boxes( filters={'area': {'min': 1e3, 'max': 1e4}}, diff --git a/model_server/extensions/chaeo/zmask.py b/model_server/extensions/chaeo/zmask.py index 557bdddb9dc763b689b65ccff34e1c6b33598124..b70afeb3939d3d9e20cd6962d29ec76eaede6d8f 100644 --- a/model_server/extensions/chaeo/zmask.py +++ b/model_server/extensions/chaeo/zmask.py @@ -104,23 +104,49 @@ class RoiSet(object): df['ebb_x1'] = (df.x1 + ebxy).apply(lambda x: min(x, w)) df['ebb_z0'] = (df.zi - ebz).apply(lambda x: max(x, 0)) df['ebb_z1'] = (df.zi + ebz).apply(lambda x: min(x, nz)) + df['ebb_h'] = df['ebb_y1'] - df['ebb_y0'] + df['ebb_w'] = df['ebb_x1'] - df['ebb_x0'] + df['ebb_nz'] = df['ebb_z1'] - df['ebb_z0'] + 1 # compute relative bounding boxes df['rel_y0'] = df.y0 - df.ebb_y0 - df['rel_y1'] = df.y1 - df.ebb_y1 + df['rel_y1'] = df.y1 - df.ebb_y0 df['rel_x0'] = df.x0 - df.ebb_x0 - df['rel_x1'] = df.x1 - df.ebb_x1 + df['rel_x1'] = df.x1 - df.ebb_x0 assert np.all(df['rel_x1'] <= (df['ebb_x1'] - df['ebb_x0'])) assert np.all(df['rel_y1'] <= (df['ebb_x1'] - df['ebb_x0'])) return df - # def get_slices(self): # TODO: actually map to DF index as new column - # sl = [] - # for ob in self.get_df().itertuples(name='LabeledObject'): - # sl.append(np.s_[ob.ebb_y0: ob.ebb_y1, ob.ebb_x0: ob.ebb_x1, :, ob.ebb_z0: ob.ebb_z1 + 1]) - # return sl + def get_slice_at(self, idx) -> tuple: + """ + Return slice object in np.s_ format that defines expanded bounding box of object + :param idx: object index (Index in DataFrame, does not necessarily start at zero) + :return: slice object + """ + ob = self.get_df().loc[idx, :].astype('int64') + return np.s_[ob.ebb_y0: ob.ebb_y1, ob.ebb_x0: ob.ebb_x1, :, ob.ebb_z0: ob.ebb_z1 + 1] + + def get_rel_slice_at(self, idx) -> tuple: + """ + Return slice object in np.s_ format that defines bounding box of an object within its expanded bounding box + :param idx: object index (Index in DataFrame, does not necessarily start at zero) + :return: slice object + """ + ob = self.get_df().loc[idx, :].astype('int64') + return np.s_[ob.rel_y0: ob.rel_y1, ob.rel_x0: ob.rel_x1, :, :] + + + def get_mask_at(self, idx) -> np.ndarray: + """ + Return 2D array describing object mask that fills (unexpanded) bounding box at index idx + :param idx: object index (Index in DataFrame, does not necessarily start at zero) + :return: np.ndarray boolean mask + """ + ob = self.get_df().loc[idx, :].astype('int64') + obmask = (self.acc_obj_ids == ob.label) + return obmask[ob.y0: ob.y1, ob.x0: ob.x1] @staticmethod def filter_df(df: pd.DataFrame, filters: RoiFilter = None) -> pd.DataFrame: @@ -170,9 +196,6 @@ class RoiSet(object): self.zmask_meta ) - def get_slices(self): - return [zm.slice for zm in self.zmask_meta] - def get_zmask(self, mask_type='boxes'): """ Return a mask of same dimensionality as raw data