Made (failing) tests of 3d product exports and TODOs needed to correct them

model_server/base/roiset.py

@@ -302,6 +302,8 @@ def make_df_from_object_ids(
 
     def _make_binary_mask(r):
         acc = InMemoryDataAccessor(acc_obj_ids.data == r.label)
+
+        # TODO: optional r.z0:z1 + 1 for case of 3d objects
         cropped = acc.get_mono(0, mip=True).crop_hw((r.y0, r.x0, (r.y1 - r.y0), (r.x1 - r.x0))).data_yx
         return cropped
 
@@ -382,6 +384,7 @@ def make_object_ids_from_df(df: pd.DataFrame, sd: dict) -> InMemoryDataAccessor:
         raise MissingSegmentationError('RoiSet dataframe does not contain segmentation')
 
     def _label_obj(r):
+        # TODO: optional r.z0:z1 + 1 for case of 3d objects
         sl = np.s_[r.y0:r.y1, r.x0:r.x1, :, r.zi:r.zi + 1]
         mask = np.expand_dims(r.binary_mask, (2, 3))
         id_mask[sl] = id_mask[sl] + r.label * mask

tests/base/test_roiset.py

@@ -248,6 +248,74 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
             self.assertEqual(h, roi.h)
             self.assertEqual(w, roi.w)
 
+class TestRoiSet3dProducts(unittest.TestCase):
+
+    where = output_path / 'run_exports_mono_3d'
+
+    def setUp(self) -> None:
+        # set up test raw data and segmentation from file
+        self.stack = generate_file_accessor(data['multichannel_zstack_raw']['path'])
+        self.stack_ch_pa = self.stack.get_mono(params['segmentation_channel'])
+        self.seg_mask_3d = generate_file_accessor(data['multichannel_zstack_mask3d']['path'])
+
+    def test_create_roiset_from_3d_obj_ids(self):
+        roiset = RoiSet.from_binary_mask(
+            self.stack_ch_pa,
+            self.seg_mask_3d,
+            allow_3d=True,
+            connect_3d=True,
+            params=RoiSetMetaParams(filters={'area': {'min': 1e2, 'max': 1e8}})
+        )
+        df = roiset.get_df()
+        df.to_csv(self.where / 'df.csv')
+
+        self.assertGreater(len(roiset.get_df()['zi'].unique()), 1)
+        self.assertTrue((df['z1'] - df['z0'] > 1).any())
+
+        # TODO: no labels are actually extending into 3D
+        roiset.acc_obj_ids.write(self.where / 'labels.tif')
+        return roiset
+
+    def test_run_export_mono_3d_patch(self):
+        p = RoiSetExportParams(**{
+            'patches_3d': {
+                'white_channel': -1,
+                'draw_bounding_box': False,
+                'draw_mask': False,
+                'pad_to': None,
+                'rgb_overlay_channels': None,
+            },
+        })
+
+        res = self.test_create_roiset_from_3d_obj_ids().run_exports(
+            self.where,
+            prefix='test',
+            params=p
+        )
+
+        # test that exported patches are padded dimension
+        for fn in res['patches_3d']:
+            pa = self.where / fn
+            self.assertTrue(pa.exists())
+            pacc = generate_file_accessor(pa)
+            self.assertEqual(pacc.chroma, 1)
+            self.assertGreater(pacc.nz, 1)
+
+    def test_run_export_mono_3d_labels_overlay(self):
+        # export via .run_exports() method
+        res = self.test_create_roiset_from_3d_obj_ids().run_exports(
+            self.where,
+            prefix='labels_overlay',
+            params=RoiSetExportParams(
+                **{
+                    'labels_overlay': {
+                        'white_channel': -1, 'transparency': 0.2, 'mip': True, 'rescale_clip': 0.00
+                    }
+                }
+            )
+        )
+        acc = generate_file_accessor(self.where / res['labels_overlay'])
+        self.assertGreater(acc.nz, 1)
 
 class TestRoiSetMultichannelProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
 
@@ -604,18 +672,6 @@ class TestRoiSetSerialization(unittest.TestCase):
         is_2d = all([self._label_is_2d(id_map.data, la) for la in labels])
         self.assertTrue(is_2d)
 
-    def test_create_roiset_from_3d_obj_ids(self):
-        id_map = get_label_ids(self.seg_mask_3d, allow_3d=True, connect_3d=False)
-        self.assertEqual(self.stack_ch_pa.shape, id_map.shape)
-
-        roiset = RoiSet.from_object_ids(
-            self.stack_ch_pa,
-            id_map,
-            params=RoiSetMetaParams(mask_type='contours')
-        )
-        self.assertEqual(roiset.count, id_map.data.max())
-        self.assertGreater(len(roiset.get_df()['zi'].unique()), 1)
-
     def test_create_roiset_from_2d_obj_ids(self):
         id_map = get_label_ids(self.seg_mask_3d, allow_3d=False)
         self.assertEqual(self.stack_ch_pa.shape[0:3], id_map.shape[0:3])