Renamed dataframe maker and moved derivative dataframe columns into separate method

model_server/base/roiset.py

@@ -129,6 +129,51 @@ def _filter_overlap_bbox(df: pd.DataFrame) -> pd.DataFrame:
     return sdf
 
 
+def _df_insert_slices(df: pd.DataFrame, shape: tuple, expand_box_by) -> pd.DataFrame:
+    h, w, c, nz = shape
+    df['h'] = df['y1'] - df['y0']
+    df['w'] = df['x1'] - df['x0']
+    ebxy, ebz = expand_box_by
+    df['ebb_y0'] = (df.y0 - ebxy).apply(lambda x: max(x, 0))
+    df['ebb_y1'] = (df.y1 + ebxy).apply(lambda x: min(x, h))
+    df['ebb_x0'] = (df.x0 - ebxy).apply(lambda x: max(x, 0))
+    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_y0
+    df['rel_x0'] = df.x0 - df.ebb_x0
+    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_y1'] - df['ebb_y0']))
+
+    df['slice'] = df.apply(
+        lambda r:
+        np.s_[int(r.y0): int(r.y1), int(r.x0): int(r.x1), :, int(r.zi): int(r.zi + 1)],
+        axis=1,
+        result_type='reduce',
+    )
+    df['expanded_slice'] = df.apply(
+        lambda r:
+        np.s_[int(r.ebb_y0): int(r.ebb_y1), int(r.ebb_x0): int(r.ebb_x1), :, int(r.ebb_z0): int(r.ebb_z1) + 1],
+        axis=1,
+        result_type='reduce',
+    )
+    df['relative_slice'] = df.apply(
+        lambda r:
+        np.s_[int(r.rel_y0): int(r.rel_y1), int(r.rel_x0): int(r.rel_x1), :, :],
+        axis=1,
+        result_type='reduce',
+    )
+    return df
+
+
 def _safe_add(a, g, b):
     assert a.dtype == b.dtype
     assert a.shape == b.shape
@@ -165,7 +210,7 @@ class RoiSet(object):
         self.params = params
 
         self._df = self.filter_df(
-            self.make_df(
+            self.make_df_from_object_ids(
                 self.acc_raw, self.acc_obj_ids, expand_box_by=params.expand_box_by
             ),
             params.filters,
@@ -203,7 +248,7 @@ class RoiSet(object):
     # call e.g. static adder
 
     @staticmethod
-    def make_df(acc_raw, acc_obj_ids, expand_box_by) -> pd.DataFrame:
+    def make_df_from_object_ids(acc_raw, acc_obj_ids, expand_box_by) -> pd.DataFrame:
         """
         Build dataframe associate object IDs with summary stats
         :param acc_raw: accessor to raw image data
@@ -231,48 +276,7 @@ class RoiSet(object):
             })
             df['zi'] = df['label'].apply(lambda x: (acc_obj_ids.data == x).sum(axis=(0, 1, 2)).argmax())
 
-        # compute expanded bounding boxes
-        h, w, c, nz = acc_raw.shape
-        df['h'] = df['y1'] - df['y0']
-        df['w'] = df['x1'] - df['x0']
-        ebxy, ebz = expand_box_by
-        df['ebb_y0'] = (df.y0 - ebxy).apply(lambda x: max(x, 0))
-        df['ebb_y1'] = (df.y1 + ebxy).apply(lambda x: min(x, h))
-        df['ebb_x0'] = (df.x0 - ebxy).apply(lambda x: max(x, 0))
-        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_y0
-        df['rel_x0'] = df.x0 - df.ebb_x0
-        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_y1'] - df['ebb_y0']))
-
-        df['slice'] = df.apply(
-            lambda r:
-            np.s_[int(r.y0): int(r.y1), int(r.x0): int(r.x1), :, int(r.zi): int(r.zi + 1)],
-            axis=1,
-            result_type='reduce',
-        )
-        df['expanded_slice'] = df.apply(
-            lambda r:
-            np.s_[int(r.ebb_y0): int(r.ebb_y1), int(r.ebb_x0): int(r.ebb_x1), :, int(r.ebb_z0): int(r.ebb_z1) + 1],
-            axis=1,
-            result_type='reduce',
-        )
-        df['relative_slice'] = df.apply(
-            lambda r:
-            np.s_[int(r.rel_y0): int(r.rel_y1), int(r.rel_x0): int(r.rel_x1), :, :],
-            axis=1,
-            result_type='reduce',
-        )
+        df = _df_insert_slices(df, acc_raw.shape, expand_box_by)
 
         # TODO: make this contingent on whether seg is included
         df['binary_mask'] = df.apply(

tests/base/test_roiset.py

@@ -71,7 +71,6 @@ class TestOverlapLogic(unittest.TestCase):
 
     def test_overlap_bbox(self):
         res = _filter_overlap_bbox(self.df)
-        print(res)
         self.assertEqual(len(res), 2)
         self.assertTrue((res.loc[0, 'overlaps_with'] == 1).all())
         self.assertTrue((res.loc[1, 'overlaps_with'] == 2).all())