Implemented and tested object identity labeling in 3D with and without...

Implemented and tested object identity labeling in 3D with and without connectivity between z-frames

model_server/base/roiset.py

@@ -65,12 +65,28 @@ class RoiSetExportParams(BaseModel):
 
 
 
-def _get_label_ids(acc_seg_mask: GenericImageDataAccessor, allow_3d=False) -> InMemoryDataAccessor:
-    if allow_3d:
+def _get_label_ids(acc_seg_mask: GenericImageDataAccessor, allow_3d=False, connect_3d=True) -> InMemoryDataAccessor:
+    """
+    Convert binary segmentation mask into either a 2D or 3D object identities map
+    :param acc_seg_mask: binary segmentation mask (mono) of either two or three dimensions
+    :param allow_3d: return a 3D map if True; return a 2D map of the mask's maximum intensity project if False
+    :param connect_3d: objects can span multiple z-positions if True; objects are unique to a single z if False
+    :return: object identities map
+    """
+    if allow_3d and connect_3d:
         nda_la = label(
             acc_seg_mask.data[:, :, 0, :]
         ).astype('uint16')
         return InMemoryDataAccessor(np.expand_dims(nda_la, 2))
+    elif allow_3d and not connect_3d:
+        nla = 0
+        la_3d = np.zeros((*acc_seg_mask.hw, 1, acc_seg_mask.nz), dtype='uint16')
+        for zi in range(0, acc_seg_mask.nz):
+            la_2d = label(acc_seg_mask.data[:, :, 0, zi]).astype('uint16')
+            la_2d[la_2d > 0] = la_2d[la_2d > 0] + nla
+            nla = la_2d.max()
+            la_3d[:, :, 0, zi] = la_2d
+        return InMemoryDataAccessor(la_3d)
     else:
         return InMemoryDataAccessor(
             label(

tests/test_roiset.py

@@ -275,19 +275,21 @@ class TestRoiSetFromZmask(unittest.TestCase):
         self.roiset = roiset
         self.zmask = InMemoryDataAccessor(roiset.get_zmask())
 
+    @staticmethod
+    def _label_is_2d(id_map, la):  # single label's zmask has same counts as its MIP
+        mask_3d = (id_map == la)
+        mask_mip = mask_3d.max(axis=-1)
+        return mask_3d.sum() == mask_mip.sum()
+
     def test_id_map_connects_z(self):
-        id_map = _get_label_ids(self.seg_mask_3d, allow_3d=True)
+        id_map = _get_label_ids(self.seg_mask_3d, allow_3d=True, connect_3d=True)
         labels = np.unique(id_map.data)[1:]
-        def _label_is_2d(la): # single label's zmask has same counts as its MIP
-            mask_3d = (id_map.data == la)
-            mask_mip = mask_3d.max(axis=-1)
-            return mask_3d.sum() == mask_mip.sum()
-        is_2d = all([_label_is_2d(la) for la in labels])
-        print([_label_is_2d(la) for la in labels])
+        is_2d = all([self._label_is_2d(id_map.data, la) for la in labels])
         self.assertFalse(is_2d)
 
-    def test_3d_zmask(self):
-        write_accessor_data_to_file(output_path / 'roiset_from_3d' / 'raw.tif', self.roiset.acc_raw)
-        write_accessor_data_to_file(output_path / 'roiset_from_3d' / 'ob_ids.tif', self.roiset.acc_obj_ids)
-        write_accessor_data_to_file(output_path / 'roiset_from_3d' / 'zmask.tif', self.zmask)
+    def test_id_map_disconnects_z(self):
+        id_map = _get_label_ids(self.seg_mask_3d, allow_3d=True, connect_3d=False)
+        labels = np.unique(id_map.data)[1:]
+        is_2d = all([self._label_is_2d(id_map.data, la) for la in labels])
+        self.assertTrue(is_2d)