More test coverage of relative bounding boxes

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}},

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