Keep all objects in dataframe and apply filter only inside zmask object loop

extensions/chaeo/examples/batch_run_patches.py

@@ -4,6 +4,7 @@ import pandas as pd
 
 from extensions.chaeo.workflows import export_patches_from_multichannel_zstack
 from extensions.ilastik.models import IlastikPixelClassifierModel
+from model_server.accessors import InMemoryDataAccessor, write_accessor_data_to_file
 
 if __name__ == '__main__':
     where_czi = Path(
@@ -43,6 +44,9 @@ if __name__ == '__main__':
         pd.DataFrame(result['timer_results'], index=[0]).to_csv(where_output / 'timer_results.csv', **csv_args)
         csv_args = {'mode': 'a', 'header': False} # append to CSV from here on
 
-
-
-
+        # export intermediate data if flagged
+        for k in result['interm'].keys():
+            write_accessor_data_to_file(
+                where_output / k / (ff.stem + '.tif'),
+                InMemoryDataAccessor(result['interm'][k])
+                )

extensions/chaeo/tests/test_zstack.py

@@ -27,7 +27,7 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         # write_accessor_data_to_file(output_path / 'obmap.tif', self.obmap)
 
     def test_zmask_makes_correct_boxes(self, mask_type='boxes', **kwargs):
-        zmask, meta, df = build_zmask_from_object_mask(
+        zmask, meta, df, interm = build_zmask_from_object_mask(
             self.obmap.get_one_channel_data(0),
             self.stack.get_one_channel_data(0),
             mask_type=mask_type,
@@ -51,6 +51,10 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         ar = meta[1]['info'].area
         self.assertGreaterEqual(sh[0] * sh[1], ar)
 
+        # assert dimensionality of intermediate data products
+        self.assertEqual(interm['label_map'].shape, zmask.shape[0:2])
+        self.assertEqual(interm['argmax'].shape, zmask.shape[0:2])
+
         return zmask, meta
 
     def test_zmask_makes_correct_contours(self):

extensions/chaeo/workflows.py

@@ -49,7 +49,7 @@ def export_patches_from_multichannel_zstack(
     ti.click('threshold_pixel_mask')
 
     # make zmask
-    zmask, zmask_meta, df = build_zmask_from_object_mask(
+    zmask, zmask_meta, df, interm = build_zmask_from_object_mask(
         obmask.get_one_channel_data(pixel_class),
         stack.get_one_channel_data(zmask_channel),
         mask_type=mask_type,
@@ -89,4 +89,5 @@ def export_patches_from_multichannel_zstack(
         'success': True,
         'timer_results': ti.events,
         'dataframe': df,
+        'interm': interm,
     }
\ No newline at end of file

extensions/chaeo/zmask.py

@@ -30,6 +30,9 @@ def build_zmask_from_object_mask(
             contour: object's contour returned by skimage.measure.find_contours
             mask: mask of object in relative frame of (optionally) expanded bounding box
         pd.DataFrame: objects, including bounding, box information after filtering
+        Dict of intermediate image products:
+            label_map: np.ndarray (h x w) where each unique object has an integer label
+            argmax: np.ndarray (h x w x 1 x 1) z-index of highest intensity in zstack
     """
 
     # validate inputs
@@ -42,7 +45,7 @@ def build_zmask_from_object_mask(
     assert zstack.hw == obmask.hw
 
     # assign object labels and build object query
-    lamap = label(obmask.data[:, :, 0, 0])
+    lamap = label(obmask.data[:, :, 0, 0]).astype('uint16')
     query_str = 'label > 0'  # always true
     if filters is not None:
         for k in filters.keys():
@@ -52,7 +55,7 @@ def build_zmask_from_object_mask(
             query_str = query_str + f' & {k} > {vmin} & {k} < {vmax}'
 
     # build dataframe of objects, assign z index to each object
-    argmax = zstack.data.argmax(axis=3, keepdims=True)[:, :, 0, 0]
+    argmax = zstack.data.argmax(axis=3, keepdims=True)[:, :, 0, 0].astype('uint16')
     df = (
         pd.DataFrame(
             regionprops_table(
@@ -61,7 +64,6 @@ def build_zmask_from_object_mask(
                 properties=('label', 'area', 'intensity_mean', 'solidity', 'bbox')
             )
         )
-        .query(query_str)
         .rename(
             columns={
                 'bbox-0': 'y0',
@@ -72,6 +74,8 @@ def build_zmask_from_object_mask(
         )
     )
     df['zi'] = df['intensity_mean'].round().astype('int')
+    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
@@ -82,7 +86,7 @@ def build_zmask_from_object_mask(
     h, w, c, nz = zstack.shape
 
     meta = []
-    for ob in df.itertuples(name='LabeledObject'):
+    for ob in df[df['keeper']].itertuples(name='LabeledObject'):
         y0 = max(ob.y0 - ebxy, 0)
         y1 = min(ob.y1 + ebxy, h - 1)
         x0 = max(ob.x0 - ebxy, 0)
@@ -133,4 +137,10 @@ def build_zmask_from_object_mask(
             sl = bb['slice']
             zi_st[sl] = 1
 
-    return zi_st, meta, df
\ No newline at end of file
+    # return intermediate image arrays
+    interm = {
+        'label_map': lamap,
+        'argmax': argmax,
+    }
+
+    return zi_st, meta, df, interm
\ No newline at end of file