Serialization/deserialization leads to identical z-masks

model_server/base/roiset.py

 from math import sqrt, floor
 from pathlib import Path
+import re
 from typing import List, Union
 from uuid import uuid4
 
@@ -17,7 +18,7 @@ from model_server.base.accessors import GenericImageDataAccessor, InMemoryDataAc
 from model_server.base.models import InstanceSegmentationModel
 from model_server.base.process import pad, rescale, resample_to_8bit, make_rgb
 from model_server.base.annotators import draw_box_on_patch, draw_contours_on_patch, draw_boxes_on_3d_image
-from model_server.base.accessors import PatchStack
+from model_server.base.accessors import generate_file_accessor, PatchStack
 from model_server.base.process import mask_largest_object
 
 
@@ -222,31 +223,6 @@ class RoiSet(object):
         return df
 
 
-    @staticmethod
-    def from_df_and_patch_masks(
-            acc_raw: GenericImageDataAccessor,
-            df: pd.DataFrame,
-            patch_masks: dict, # dict of ndarray, where key is integer label
-            params: RoiSetMetaParams = RoiSetMetaParams(),
-    ):
-        assert len(df) == len(patch_masks)
-        zmask = np.zeros((*acc_raw.hw, 1, acc_raw.nz), dtype=bool)
-
-        def _label_obj(r):
-            # sl = np.s_[r.ebb_y0:r.ebb_y1, r.ebb_x0:r.ebb_x1, :, r.zi:r.zi + 1]
-            sl = np.s_[r.y0:r.y1, r.x0:r.x1, :, r.zi:r.zi + 1]
-            rsl = np.s_[r.rel_y0:r.rel_y1, r.rel_x0:r.rel_x1, :, :]
-            try:
-                zmask[sl] = patch_masks[r.label][rsl]
-                print(f'ebb (h, w) = {r.ebb_h, r.ebb_w}')
-            except ValueError as ve:
-                print(ve)
-
-        df.apply(lambda x: _label_obj(x), axis=1)
-        obj_ids = _get_label_ids(zmask, allow_3d=True, connect_3d=False)
-        return RoiSet(acc_raw, obj_ids, params)
-
-
     @staticmethod
     def filter_df(df: pd.DataFrame, filters: RoiFilter = None) -> pd.DataFrame:
         query_str = 'label > 0'  # always true
@@ -619,6 +595,57 @@ class RoiSet(object):
         )
         return record
 
+    @staticmethod
+    def deserialize(acc_raw: GenericImageDataAccessor, where: Path, prefix=''):
+        df = pd.read_csv(where / 'dataframe' / (prefix + '.csv'))[['label', 'zi', 'y0', 'y1', 'x0', 'x1']]
+
+        id_mask = np.zeros((*acc_raw.hw, 1, acc_raw.nz), dtype='uint16')
+
+        for f in (where / 'tight_patch_masks').iterdir():
+            acc = generate_file_accessor(f)
+            la = int(re.search(r'la([\d]+)', str(f)).group(1))
+            zi = int(re.search(r'zi([\d]+)', str(f)).group(1))
+            roi_q = df.loc[(df.label == la) & (df.zi == zi), :]
+            assert len(roi_q) == 1, f'Did not find exactly one corresponding ROI for patch mask {f}'
+
+        def _label_obj(r):
+            sl = np.s_[r.y0:r.y1, r.x0:r.x1, :, r.zi:r.zi + 1]
+            ext = 'png'
+            fname = f'{prefix}-la{r.label:04d}-zi{r.zi:04d}.{ext}'
+            try:
+                ma_acc = generate_file_accessor(where / 'tight_patch_masks' / fname)
+                id_mask[sl] = r.label * ma_acc.data
+            except Exception as e:
+                raise DeserializeRoiSet(e)
+
+        df.apply(_label_obj, axis=1)
+        return RoiSet(acc_raw, InMemoryDataAccessor(id_mask))
+
+
+    @staticmethod
+    def from_df_and_patch_masks(
+            acc_raw: GenericImageDataAccessor,
+            df: pd.DataFrame,
+            patch_masks: dict, # dict of ndarray, where key is integer label
+            params: RoiSetMetaParams = RoiSetMetaParams(),
+    ):
+        assert len(df) == len(patch_masks)
+        zmask = np.zeros((*acc_raw.hw, 1, acc_raw.nz), dtype=bool)
+
+        def _label_obj(r):
+            # sl = np.s_[r.ebb_y0:r.ebb_y1, r.ebb_x0:r.ebb_x1, :, r.zi:r.zi + 1]
+            sl = np.s_[r.y0:r.y1, r.x0:r.x1, :, r.zi:r.zi + 1]
+            rsl = np.s_[r.rel_y0:r.rel_y1, r.rel_x0:r.rel_x1, :, :]
+            try:
+                zmask[sl] = patch_masks[r.label][rsl]
+                print(f'ebb (h, w) = {r.ebb_h, r.ebb_w}')
+            except ValueError as ve:
+                print(ve)
+
+        df.apply(lambda x: _label_obj(x), axis=1)
+        obj_ids = _get_label_ids(zmask, allow_3d=True, connect_3d=False)
+        return RoiSet(acc_raw, obj_ids, params)
+
 
 def project_stack_from_focal_points(
         xx: np.ndarray,
@@ -667,3 +694,9 @@ def project_stack_from_focal_points(
     )
 
 
+
+class Error(Exception):
+    pass
+
+class DeserializeRoiSet(Error):
+    pass
\ No newline at end of file

tests/test_roiset.py

@@ -381,12 +381,8 @@ class TestRoiSetFromZmask(unittest.TestCase):
 
     def test_create_roiset_from_df_and_patch_masks(self):
         ref_roiset = self.test_create_roiset_from_2d_obj_ids()
-        res = ref_roiset.run_exports(
-            output_path / 'roiset_from_3d',
-            roiset_test_data['pipeline_params']['segmentation_channel'],
-            'ref',
-            params=RoiSetExportParams()
-        )
+        where_ser = output_path / 'roiset_from_3d'
+        ref_roiset.serialize(where_ser, prefix='ref')
         where_df = output_path / 'roiset_from_3d' / 'dataframe' / 'ref.csv'
         self.assertTrue(where_df.exists())
         df_test = pd.read_csv(where_df)
@@ -403,21 +399,6 @@ class TestRoiSetFromZmask(unittest.TestCase):
             self.assertEqual((roi.h, roi.w), m_acc.hw)
 
         # make another RoiSet from just the data table, raw images, and (tight) patch masks
-        df_test = pd.read_csv(where_df)
-        zmask = np.zeros((*self.stack.hw, 1, self.stack.nz), dtype=bool)
-
-        # fn = output_path / 'roiset_from_3d' / 'patch_masks' / 'ref-la{:04d}-zi{:04d}.png'
-        # patch_masks = {}
-        #
-        # def _label_obj(r):
-        #     sl = np.s_[r.ebb_y0:r.ebb_y1, r.ebb_x0:r.ebb_x1, :, r.zi:r.zi + 1]
-        #     self.assertEqual(str(sl), r.slice)
-        #     patch_masks[r.label] = generate_file_accessor(str(fn).format(r.label, r.zi)).data
-        #     zmask[sl] = True
-        #
-        # df_test.apply(lambda x: _label_obj(x), axis=1)
-        #
-        # roiset_test = RoiSet.from_df_and_patch_masks(self.stack, df_test, patch_masks)
-        # print('')
-
+        test_roiset = RoiSet.deserialize(self.stack, where_ser, prefix='ref')
+        self.assertTrue(np.all(ref_roiset.get_zmask() == test_roiset.get_zmask()))