Merge branch 'dev_obj_det' into 'staging'

RoiSet facilitates object detection models

See merge request rhodes/model_server!57

model_server/base/accessors.py

@@ -49,13 +49,51 @@ class GenericImageDataAccessor(ABC):
             nda = self.data.take(indices=carr, axis=self._ga('C'))
             return self._derived_accessor(nda)
 
+
+    def get_zi(self, zi: int):
+        """
+        Return a new accessor of a specific z-coordinate
+        """
+        return self._derived_accessor(
+            self.data.take(
+                indices=[zi],
+                axis=self._ga('Z')
+            )
+        )
+
+    def get_mip(self):
+        """
+        Return a new accessor of maximum intensity projection (MIP) along z-axis
+        """
+        return self.apply(lambda x: x.max(axis=self._ga('Z'), keepdims=True))
+
     def get_mono(self, channel: int, mip: bool = False):
         return self.get_channels([channel], mip=mip)
 
+    def get_z_argmax(self):
+        return self.apply(lambda x: x.argmax(axis=self.get_axis('Z')))
+
+    def get_focus_vector(self):
+        return self.data.sum(axis=(0, 1, 2))
+
+    @property
+    def data_xy(self) -> np.ndarray:
+        if not self.chroma == 1 and self.nz == 1:
+            raise InvalidDataShape('Can only return XY array from accessors with a single channel and single z-level')
+        else:
+            return self.data[:, :, 0, 0]
+
+    @property
+    def data_xyz(self) -> np.ndarray:
+        if not self.chroma == 1:
+            raise InvalidDataShape('Can only return XYZ array from accessors with a single channel')
+        else:
+            return self.data[:, :, 0, :]
+
     def _gc(self, channels):
         return self.get_channels(list(channels))
 
-    def _unique(self):
+    def unique(self):
         return np.unique(self.data, return_counts=True)
 
     @property
@@ -75,6 +113,15 @@ class GenericImageDataAccessor(ABC):
     def _ga(self, arg):
         return self.get_axis(arg)
 
+    def crop_hw(self, yxhw: tuple):
+        """
+        Return subset of data cropped in X and Y
+        :param yxhw: tuple (Y, X, H, W)
+        :return: InMemoryDataAccessor of size (H x W), starting at (Y, X)
+        """
+        y, x, h, w = yxhw
+        return InMemoryDataAccessor(self.data[y: (y + h), x: (x + w), :, :])
+
     @property
     def hw(self):
         """
@@ -120,6 +167,14 @@ class GenericImageDataAccessor(ABC):
             func(self.data)
         )
 
+    @property
+    def info(self):
+        return {
+            'shape_dict': self.shape_dict,
+            'dtype': str(self.dtype),
+            'filepath': '',
+        }
+
 class InMemoryDataAccessor(GenericImageDataAccessor):
     def __init__(self, data):
         self._data = self.conform_data(data)
@@ -139,6 +194,12 @@ class GenericImageFileAccessor(GenericImageDataAccessor): # image data is loaded
     def read(fp: Path):
         return generate_file_accessor(fp)
 
+    @property
+    def info(self):
+        d = super().info
+        d['filepath'] = self.fpath.__str__()
+        return d
+
 class TifSingleSeriesFileAccessor(GenericImageFileAccessor):
     def __init__(self, fpath: Path):
         super().__init__(fpath)
@@ -240,7 +301,7 @@ class CziImageFileAccessor(GenericImageFileAccessor):
 
 def write_accessor_data_to_file(fpath: Path, acc: GenericImageDataAccessor, mkdir=True) -> bool:
     """
-    Export an image accessor to file.
+    Export an image accessor to file
     :param fpath: complete path including filename and extension
     :param acc: image accessor to be written
     :param mkdir: create any needed subdirectories in fpath if True
@@ -287,7 +348,7 @@ def write_accessor_data_to_file(fpath: Path, acc: GenericImageDataAccessor, mkdi
 def generate_file_accessor(fpath):
     """
     Given an image file path, return an image accessor, assuming the file is a supported format and represents
-    a single position array, which may be single or multi-channel, single plane or z-stack.
+    a single position array, which may be single or multichannel, single plane or z-stack.
     """
     if str(fpath).upper().endswith('.TIF') or str(fpath).upper().endswith('.TIFF'):
         return TifSingleSeriesFileAccessor(fpath)
@@ -379,6 +440,11 @@ class PatchStack(InMemoryDataAccessor):
         else:
             tifffile.imwrite(fpath, tzcyx, imagej=True)
 
+    def write(self, fp: Path, mkdir=True):
+        if mkdir:
+            fp.parent.mkdir(parents=True, exist_ok=True)
+        self.export_pyxcz(fp)
+
     @property
     def shape_dict(self):
         return dict(zip(('P', 'Y', 'X', 'C', 'Z'), self.data.shape))
@@ -437,7 +503,6 @@ def make_patch_stack_from_file(fpath):  # interpret t-dimension as patch positio
     return PatchStack(pyxcz)
 
 
-
 class Error(Exception):
     pass
 

model_server/base/process.py

@@ -18,7 +18,7 @@ def is_mask(img):
         return True
     elif img.dtype == 'uint8':
         unique = np.unique(img)
-        if unique.shape[0] == 2 and np.all(unique == [0, 255]):
+        if unique.shape[0] <= 2 and np.all(unique == [0, 255]):
             return True
     return False
 
@@ -136,7 +136,14 @@ def smooth(img: np.ndarray, sig: float) -> np.ndarray:
     :param sig: threshold parameter
     :return: smoothed image
     """
-    return gaussian(img, sig)
+    ga = gaussian(img, sig, preserve_range=True)
+    if is_mask(img):
+        if img.dtype == 'bool':
+            return ga > ga.mean()
+        elif img.dtype == 'uint8':
+            return (255 * (ga > ga.mean())).astype('uint8')
+    else:
+        return ga
 
 class Error(Exception):
     pass

model_server/base/util.py

@@ -8,6 +8,7 @@ import pandas as pd
 
 from .accessors import InMemoryDataAccessor, write_accessor_data_to_file
 from .models import Model
+from .roiset import filter_df_overlap_seg, RoiSet
 
 def autonumber_new_directory(where: str, prefix: str) -> str:
     """
@@ -163,4 +164,4 @@ def loop_workflow(
                 )
 
     if len(failures) > 0:
-        pd.DataFrame(failures).to_csv(Path(output_folder_path) / 'failures.csv')
+        pd.DataFrame(failures).to_csv(Path(output_folder_path) / 'failures.csv')
\ No newline at end of file

tests/base/test_accessors.py

@@ -61,6 +61,52 @@ class TestCziImageFileAccess(unittest.TestCase):
         sc = cf.get_mono(c, mip=True)
         self.assertEqual(sc.shape, (h, w, 1, 1))
 
+    def test_get_single_channel_argmax_from_zstack(self):
+        w = 256
+        h = 512
+        nc = 4
+        nz = 11
+        c = 3
+        cf = InMemoryDataAccessor(np.random.rand(h, w, nc, nz))
+        am = cf.get_mono(c).get_z_argmax()
+        self.assertEqual(am.shape, (h, w, 1, 1))
+        self.assertTrue(np.all(am.unique()[0] == range(0, nz)))
+
+    def test_get_single_channel_z_series_from_zstack(self):
+        w = 256
+        h = 512
+        nc = 4
+        nz = 11
+        c = 3
+        cf = InMemoryDataAccessor(np.random.rand(h, w, nc, nz))
+        zs = cf.get_mono(c).get_focus_vector()
+        self.assertEqual(zs.shape, (nz, ))
+
+    def test_get_zi(self):
+        w = 256
+        h = 512
+        nc = 4
+        nz = 11
+        zi = 5
+        cf = InMemoryDataAccessor(_random_int(h, w, nc, nz))
+        sz = cf.get_zi(zi)
+        self.assertEqual(sz.shape_dict['Z'], 1)
+
+        self.assertTrue(np.all(sz.data[:, :, :, 0] == cf.data[:, :, :, zi]))
+
+    def test_crop_yx(self):
+        w = 256
+        h = 512
+        nc = 4
+        nz = 11
+        cf = InMemoryDataAccessor(_random_int(h, w, nc, nz))
+
+        yxhw = (100, 200, 10, 20)
+        sc = cf.crop_hw(yxhw)
+        self.assertEqual(sc.shape_dict['Z'], nz)
+        self.assertEqual(sc.shape_dict['C'], nc)
+        self.assertEqual(sc.hw, yxhw[2:])
+
     def test_write_single_channel_tif(self):
         ch = 4
         cf = CziImageFileAccessor(data['czifile']['path'])

tests/base/test_roiset.py

@@ -6,16 +6,19 @@ from pathlib import Path
 
 import pandas as pd
 
-from model_server.base.roiset import RoiSetExportParams, RoiSetMetaParams
+from model_server.base.roiset import filter_df_overlap_bbox, filter_df_overlap_seg, RoiSetExportParams, RoiSetMetaParams
 from model_server.base.roiset import RoiSet
 from model_server.base.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file, PatchStack
 from model_server.base.models import DummyInstanceSegmentationModel
+from model_server.base.process import smooth
 import model_server.conf.testing as conf
 
 data = conf.meta['image_files']
 output_path = conf.meta['output_path']
 params = conf.meta['roiset']
 
+
+
 class BaseTestRoiSetMonoProducts(object):
 
     def setUp(self) -> None:
@@ -28,7 +31,7 @@ class BaseTestRoiSetMonoProducts(object):
 class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
 
     def _make_roi_set(self, mask_type='boxes', **kwargs):
-        roiset = RoiSet.from_segmentation(
+        roiset = RoiSet.from_binary_mask(
             self.stack_ch_pa,
             self.seg_mask,
             params=RoiSetMetaParams(
@@ -69,7 +72,7 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
         acc_zstack_slice = InMemoryDataAccessor(self.stack_ch_pa.data[:, :, :, 0])
         self.assertEqual(acc_zstack_slice.nz, 1)
 
-        roiset = RoiSet.from_segmentation(acc_zstack_slice, self.seg_mask, params=RoiSetMetaParams(mask_type='boxes'))
+        roiset = RoiSet.from_binary_mask(acc_zstack_slice, self.seg_mask, params=RoiSetMetaParams(mask_type='boxes'))
         zmask = roiset.get_zmask()
 
         zmask_acc = InMemoryDataAccessor(zmask)
@@ -77,7 +80,7 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
 
     def test_create_roiset_with_no_objects(self):
         zero_obmap = InMemoryDataAccessor(np.zeros(self.seg_mask.shape, self.seg_mask.dtype))
-        roiset = RoiSet(self.stack_ch_pa, zero_obmap)
+        roiset = RoiSet.from_object_ids(self.stack_ch_pa, zero_obmap)
         self.assertEqual(roiset.count, 0)
 
     def test_slices_are_valid(self):
@@ -162,26 +165,6 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
         result = generate_file_accessor(where / file)
         self.assertEqual(result.shape, roiset.acc_raw.shape)
 
-    def test_flatten_image(self):
-        roiset = RoiSet.from_segmentation(self.stack_ch_pa, self.seg_mask, params=RoiSetMetaParams(mask_type='boxes'))
-        df = roiset.get_df()
-
-        from model_server.base.roiset import project_stack_from_focal_points
-
-        img = project_stack_from_focal_points(
-            df['centroid-0'].to_numpy(),
-            df['centroid-1'].to_numpy(),
-            df['zi'].to_numpy(),
-            self.stack,
-            degree=4,
-        )
-
-        self.assertEqual(img.shape[0:2], self.stack.shape[0:2])
-
-        write_accessor_data_to_file(
-            output_path / 'flattened.tif',
-            InMemoryDataAccessor(img)
-        )
 
     def test_make_binary_masks(self):
         roiset = self._make_roi_set()
@@ -200,26 +183,51 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
         roiset = self._make_roi_set()
         roiset.classify_by('dummy_class', [0], DummyInstanceSegmentationModel())
         self.assertTrue(all(roiset.get_df()['classify_by_dummy_class'].unique() == [1]))
-        self.assertTrue(all(np.unique(roiset.object_class_maps['dummy_class'].data) == [0, 1]))
+        self.assertTrue(all(np.unique(roiset.get_object_class_map('dummy_class').data) == [0, 1]))
         return roiset
 
     def test_classify_by_multiple_channels(self):
-        roiset = RoiSet.from_segmentation(self.stack, self.seg_mask)
+        roiset = RoiSet.from_binary_mask(self.stack, self.seg_mask, params=RoiSetMetaParams(deproject_channel=0))
         roiset.classify_by('dummy_class', [0, 1], DummyInstanceSegmentationModel())
         self.assertTrue(all(roiset.get_df()['classify_by_dummy_class'].unique() == [1]))
-        self.assertTrue(all(np.unique(roiset.object_class_maps['dummy_class'].data) == [0, 1]))
+        self.assertTrue(all(np.unique(roiset.get_object_class_map('dummy_class').data) == [0, 1]))
         return roiset
 
+    def test_transfer_classification(self):
+        roiset1 = RoiSet.from_binary_mask(self.stack, self.seg_mask, params=RoiSetMetaParams(deproject_channel=0))
+
+        # prepare alternative mask and compare
+        smoothed_mask = self.seg_mask.apply(lambda x: smooth(x, sig=1.5))
+        roiset2 = RoiSet.from_binary_mask(self.stack, smoothed_mask, params=RoiSetMetaParams(deproject_channel=0))
+        dmask = (self.seg_mask.data / 255) + (smoothed_mask.data / 255)
+        self.assertTrue(np.all(np.unique(dmask) == [0, 1, 2]))
+        total_iou = (dmask == 2).sum() / ((dmask == 1).sum() + (dmask == 2).sum())
+        self.assertGreater(total_iou, 0.6)
+
+        # classify first RoiSet
+        roiset1.classify_by('dummy_class', [0, 1], DummyInstanceSegmentationModel())
+
+        self.assertTrue('dummy_class' in roiset1.classification_columns)
+        self.assertFalse('dummy_class' in roiset2.classification_columns)
+        res = roiset2.get_instance_classification(roiset1)
+        self.assertTrue('dummy_class' in roiset2.classification_columns)
+        self.assertLess(
+            roiset2.get_df().classify_by_dummy_class.count(),
+            roiset1.get_df().classify_by_dummy_class.count(),
+        )
+
+
     def test_classify_by_with_derived_channel(self):
         class ModelWithDerivedInputs(DummyInstanceSegmentationModel):
             def infer(self, img, mask):
                 return PatchStack(super().infer(img, mask).data * img.chroma)
 
-        roiset = RoiSet.from_segmentation(
+        roiset = RoiSet.from_binary_mask(
             self.stack,
             self.seg_mask,
             params=RoiSetMetaParams(
                 filters={'area': {'min': 1e3, 'max': 1e4}},
+                deproject_channel=0,
             )
         )
         roiset.classify_by(
@@ -232,7 +240,7 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
             ]
         )
         self.assertTrue(all(roiset.get_df()['classify_by_multiple_input_model'].unique() == [4]))
-        self.assertTrue(all(np.unique(roiset.object_class_maps['multiple_input_model'].data) == [0, 4]))
+        self.assertTrue(all(np.unique(roiset.get_object_class_map('multiple_input_model').data) == [0, 4]))
 
         self.assertEqual(len(roiset.accs_derived), 2)
         for di in roiset.accs_derived:
@@ -281,13 +289,14 @@ class TestRoiSetMultichannelProducts(BaseTestRoiSetMonoProducts, unittest.TestCa
 
     def setUp(self) -> None:
         super().setUp()
-        self.roiset = RoiSet.from_segmentation(
+        self.roiset = RoiSet.from_binary_mask(
             self.stack,
             self.seg_mask,
             params=RoiSetMetaParams(
                 expand_box_by=(128, 2),
                 mask_type='boxes',
                 filters={'area': {'min': 1e3, 'max': 1e4}},
+                deproject_channel=0,
             )
         )
 
@@ -520,7 +529,7 @@ class TestRoiSetMultichannelProducts(BaseTestRoiSetMonoProducts, unittest.TestCa
             self.assertEqual(pacc.chroma, 1)
 
 
-from model_server.base.roiset import _get_label_ids
+from model_server.base.roiset import get_label_ids
 class TestRoiSetSerialization(unittest.TestCase):
 
     def setUp(self) -> None:
@@ -528,6 +537,7 @@ class TestRoiSetSerialization(unittest.TestCase):
         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'])
+        self.seg_mask_2d = generate_file_accessor(data['multichannel_zstack_raw']['path'])
 
     @staticmethod
     def _label_is_2d(id_map, la):  # single label's zmask has same counts as its MIP
@@ -536,22 +546,22 @@ class TestRoiSetSerialization(unittest.TestCase):
         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, connect_3d=True)
+        id_map = get_label_ids(self.seg_mask_3d, allow_3d=True, connect_3d=True)
         labels = np.unique(id_map.data)[1:]
         is_2d = all([self._label_is_2d(id_map.data, la) for la in labels])
         self.assertFalse(is_2d)
 
     def test_id_map_disconnects_z(self):
-        id_map = _get_label_ids(self.seg_mask_3d, allow_3d=True, connect_3d=False)
+        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)
 
     def test_create_roiset_from_3d_obj_ids(self):
-        id_map = _get_label_ids(self.seg_mask_3d, allow_3d=True, connect_3d=False)
+        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(
+        roiset = RoiSet.from_object_ids(
             self.stack_ch_pa,
             id_map,
             params=RoiSetMetaParams(mask_type='contours')
@@ -560,11 +570,11 @@ class TestRoiSetSerialization(unittest.TestCase):
         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)
+        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])
         self.assertEqual(id_map.nz, 1)
 
-        roiset = RoiSet(
+        roiset = RoiSet.from_object_ids(
             self.stack_ch_pa,
             id_map,
             params=RoiSetMetaParams(mask_type='contours')
@@ -593,6 +603,7 @@ class TestRoiSetSerialization(unittest.TestCase):
             m_acc = generate_file_accessor(pmf)
             self.assertEqual((roi.h, roi.w), m_acc.hw)
             patch_filenames.append(pmf.name)
+            self.assertEqual(m_acc.nz, 1)
 
         # make another RoiSet from just the data table, raw images, and (tight) patch masks
         test_roiset = RoiSet.deserialize(self.stack_ch_pa, where_ser, prefix='ref')
@@ -614,3 +625,191 @@ class TestRoiSetSerialization(unittest.TestCase):
             t_acc = generate_file_accessor(pt)
             self.assertTrue(np.all(r_acc.data == t_acc.data))
 
+
+class TestRoiSetObjectDetection(unittest.TestCase):
+
+    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_bounding_boxes(self):
+        from skimage.measure import label, regionprops, regionprops_table
+
+        mask = self.seg_mask_3d
+        labels = label(mask.data_xyz, connectivity=3)
+        table = pd.DataFrame(
+            regionprops_table(labels)
+        ).rename(
+            columns={'bbox-0': 'y', 'bbox-1': 'x', 'bbox-2': 'zi', 'bbox-3': 'y1', 'bbox-4': 'x1'}
+        ).drop(
+            columns=['bbox-5']
+        )
+        table['w'] = table['x1'] - table['x']
+        table['h'] = table['y1'] - table['y']
+        bboxes = table[['y', 'x', 'h', 'w']].to_dict(orient='records')
+
+        roiset_bbox = RoiSet.from_bounding_boxes(self.stack_ch_pa, bboxes)
+        self.assertTrue('label' in roiset_bbox.get_df().columns)
+        patches_bbox = roiset_bbox.get_patches_acc()
+        self.assertEqual(len(table), patches_bbox.count)
+
+
+        # roiset w/ seg for comparison
+        roiset_seg = RoiSet.from_binary_mask(self.stack_ch_pa, mask, allow_3d=True)
+        patches_seg = roiset_seg.get_patches_acc()
+
+        # test bounding box dimensions match those from RoiSet generated directly from segmentation
+        self.assertEqual(roiset_seg.count, roiset_bbox.count)
+        for i in range(0, roiset_seg.count):
+            self.assertEqual(patches_seg.iat(0, crop=True).shape, patches_bbox.iat(0, crop=True).shape)
+
+        # test that serialization does not write patch masks
+        roiset_ser_path = output_path / 'roiset_from_bbox'
+        dd = roiset_bbox.serialize(roiset_ser_path)
+        self.assertTrue('tight_patch_masks' not in dd.keys())
+        self.assertFalse((roiset_ser_path / 'tight_patch_masks').exists())
+
+        # test that deserialized RoiSet matches the original
+        roiset_des = RoiSet.deserialize(self.stack_ch_pa, roiset_ser_path)
+        self.assertEqual(roiset_des.count, roiset_bbox.count)
+        for i in range(0, roiset_des.count):
+            self.assertEqual(patches_seg.iat(0, crop=True).shape, patches_bbox.iat(0, crop=True).shape)
+        self.assertTrue((roiset_bbox.get_zmask() == roiset_des.get_zmask()).all())
+
+
+class TestRoiSetPolygons(BaseTestRoiSetMonoProducts, unittest.TestCase):
+
+    def test_compute_polygons(self):
+        roiset_ref = RoiSet.from_binary_mask(
+            self.stack_ch_pa,
+            self.seg_mask,
+            params=RoiSetMetaParams(
+                mask_type='contours',
+                filters={'area': {'min': 1e1, 'max': 1e6}}
+            )
+        )
+
+        poly = roiset_ref.get_polygons()
+        roiset_test = RoiSet.from_polygons_2d(self.stack_ch_pa, poly)
+        binary_poly = (roiset_test.acc_obj_ids.get_mono(0, mip=True).data > 0)
+        self.assertEqual(self.seg_mask.shape, binary_poly.shape)
+
+
+        # most mask pixels are within in fitted polygon
+        test_mask = np.logical_and(
+            np.logical_not(binary_poly),
+            (self.seg_mask.data == 255)
+        )
+        self.assertLess(test_mask.sum() / test_mask.size, 0.001)
+
+        # output results
+        od = output_path / 'polygons'
+        write_accessor_data_to_file(od / 'from_polygons.tif', InMemoryDataAccessor(binary_poly))
+        write_accessor_data_to_file(od / 'ref_mask.tif', self.seg_mask)
+        write_accessor_data_to_file(od / 'diff.tif', InMemoryDataAccessor(test_mask))
+
+
+    def test_overlap_bbox(self):
+        df = pd.DataFrame({
+            'x0': [0, 1, 2, 1, 1],
+            'x1': [2, 3, 4, 3, 3],
+            'y0': [0, 0, 0, 2, 0],
+            'y1': [2, 2, 2, 3, 2],
+            'zi': [0, 0, 0, 0, 1],
+        })
+
+        res = filter_df_overlap_bbox(df)
+        self.assertEqual(len(res), 4)
+        self.assertTrue((res.loc[0, 'overlaps_with'] == [1]).all())
+        self.assertTrue((res.loc[1, 'overlaps_with'] == [0, 2]).all())
+        self.assertTrue((res.bbox_intersec == 2).all())
+        return res
+
+
+    def test_overlap_bbox_multiple(self):
+        df1 = pd.DataFrame({
+            'x0': [0, 1],
+            'x1': [2, 3],
+            'y0': [0, 0],
+            'y1': [2, 2],
+            'zi': [0, 0],
+        })
+        df2 = pd.DataFrame({
+            'x0': [2],
+            'x1': [4],
+            'y0': [0],
+            'y1': [2],
+            'zi': [0],
+        })
+        res = filter_df_overlap_bbox(df1, df2)
+        self.assertTrue((res.loc[1, 'overlaps_with'] == [0]).all())
+        self.assertEqual(len(res), 1)
+        self.assertTrue((res.bbox_intersec == 2).all())
+
+
+    def test_overlap_seg(self):
+        df = pd.DataFrame({
+            'x0': [0, 1, 2],
+            'x1': [2, 3, 4],
+            'y0': [0, 0, 0],
+            'y1': [2, 2, 2],
+            'zi': [0, 0, 0],
+            'binary_mask': [
+                [
+                    [1, 1],
+                    [1, 0]
+                ],
+                [
+                    [0, 1],
+                    [1, 1]
+                ],
+                [
+                    [1, 1],
+                    [1, 1]
+                ],
+            ]
+        })
+
+        res = filter_df_overlap_seg(df)
+        self.assertTrue((res.loc[res.seg_overlaps, :].index == [1, 2]).all())
+        self.assertTrue((res.loc[res.seg_overlaps, 'seg_iou'] == 0.4).all())
+
+    def test_overlap_seg_multiple(self):
+        df1 = pd.DataFrame({
+            'x0': [0, 1],
+            'x1': [2, 3],
+            'y0': [0, 0],
+            'y1': [2, 2],
+            'zi': [0, 0],
+            'binary_mask': [
+                [
+                    [1, 1],
+                    [1, 0]
+                ],
+                [
+                    [0, 1],
+                    [1, 1]
+                ],
+            ]
+        })
+        df2 = pd.DataFrame({
+            'x0': [2],
+            'x1': [4],
+            'y0': [0],
+            'y1': [2],
+            'zi': [0],
+            'binary_mask': [
+                [
+                    [1, 1],
+                    [1, 1]
+                ],
+            ]
+        })
+        res = filter_df_overlap_seg(df1, df2)
+        self.assertTrue((res.loc[1, 'overlaps_with'] == [0]).all())
+        self.assertEqual(len(res), 1)
+        self.assertTrue((res.bbox_intersec == 2).all())
+        self.assertTrue((res.loc[res.seg_overlaps, :].index == [1]).all())
+        self.assertTrue((res.loc[res.seg_overlaps, 'seg_iou'] == 0.4).all())

tests/test_ilastik/test_ilastik.py

@@ -5,7 +5,7 @@ import numpy as np
 from model_server.base.accessors import CziImageFileAccessor, generate_file_accessor, InMemoryDataAccessor, PatchStack, write_accessor_data_to_file
 from model_server.extensions.ilastik import models as ilm
 from model_server.extensions.ilastik.workflows import infer_px_then_ob_model
-from model_server.base.roiset import _get_label_ids, RoiSet, RoiSetMetaParams
+from model_server.base.roiset import get_label_ids, RoiSet, RoiSetMetaParams
 from model_server.base.workflows import classify_pixels
 import model_server.conf.testing as conf
 
@@ -363,7 +363,7 @@ class TestIlastikOnMultichannelInputs(conf.TestServerBaseClass):
         acc_input = generate_file_accessor(self.pa_input_image)
         acc_obmap = generate_file_accessor(res.object_map_filepath)
         self.assertEqual(acc_obmap.hw, acc_input.hw)
-        self.assertEqual(len(acc_obmap._unique()[1]), 3)
+        self.assertEqual(len(acc_obmap.unique()[1]), 3)
 
 
     def test_api(self):
@@ -401,9 +401,9 @@ class TestIlastikObjectClassification(unittest.TestCase):
         stack_ch_pa = stack.get_mono(conf.meta['roiset']['patches_channel'])
         seg_mask = generate_file_accessor(data['multichannel_zstack_mask2d']['path'])
 
-        self.roiset = RoiSet(
+        self.roiset = RoiSet.from_binary_mask(
             stack_ch_pa,
-            _get_label_ids(seg_mask),
+            seg_mask,
             params=RoiSetMetaParams(
                 mask_type='boxes',
                 filters={'area': {'min': 1e3, 'max': 1e4}},