RoiSet method to create an aggregate column based on dataframe query

model_server/base/roiset.py

@@ -611,6 +611,25 @@ class RoiSet(object):
             se[roi.Index] = oc
         self.set_classification(f'classify_by_{name}', se)
 
+    def aggregate_classifications(self, query: str, name: str = 'aggregate'):
+        """
+        Run query on DataFrame and put the results in a new boolean column
+        """
+        cname = 'classify_by_' + name
+        if cname in self._df.columns:
+            raise DataFrameQueryError(f'Name {cname} is already used in RoiSet dataframe')
+
+        if self.count == 0:
+            self._df[cname] = None
+            return True
+
+        try:
+            self.set_classification(
+                cname,
+                self._df.eval(query)
+            )
+        except Exception as e:
+            raise DataFrameQueryError(e)
 
     def get_instance_classification(self, roiset_from: Self, iou_min: float = 0.5) -> pd.DataFrame:
         """
@@ -1226,6 +1245,10 @@ class BoundingBoxError(Error):
     pass
 
 
+class DataFrameQueryError(Error):
+    pass
+
+
 class DeserializeRoiSetError(Error):
     pass
 

tests/base/test_roiset.py

@@ -7,7 +7,7 @@ from pathlib import Path
 import pandas as pd
 
 from model_server.base.process import smooth
-from model_server.base.roiset import filter_df_overlap_bbox, filter_df_overlap_seg, RoiSet, RoiSetExportParams, RoiSetMetaParams
+from model_server.base.roiset import filter_df_overlap_bbox, filter_df_overlap_seg, IntensityThresholdInstanceMaskSegmentationModel, RoiSet, RoiSetExportParams, RoiSetMetaParams
 from model_server.base.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
 import model_server.conf.testing as conf
 from model_server.conf.testing import DummyInstanceMaskSegmentationModel
@@ -820,8 +820,7 @@ class TestRoiSetPolygons(BaseTestRoiSetMonoProducts, unittest.TestCase):
         self.assertTrue((res.loc[res.seg_overlaps, 'seg_iou'] == 0.4).all())
 
 class TestIntensityThresholdObjectModel(BaseTestRoiSetMonoProducts, unittest.TestCase):
-    def test_permissive_instance_segmentation(self):
-        from model_server.base.roiset import IntensityThresholdInstanceMaskSegmentationModel
+    def test_instance_segmentation(self):
 
         img = self.stack.get_mono(channel=0, mip=True)
         mask = self.seg_mask
@@ -833,4 +832,39 @@ class TestIntensityThresholdObjectModel(BaseTestRoiSetMonoProducts, unittest.Tes
         mask.write(output_path / 'TestIntensityThresholdObjectModel' / 'mask.tif')
         obmap.write(output_path / 'TestIntensityThresholdObjectModel' / 'obmap.tif')
 
-        self.assertGreater((mask.data > 0).sum(), (obmap.data > 0).sum())
\ No newline at end of file
+        self.assertGreater((mask.data > 0).sum(), (obmap.data > 0).sum())
+
+    def test_roiset_with_instance_segmentation(self):
+        roiset = RoiSet.from_binary_mask(
+            self.stack,
+            self.seg_mask,
+            params=RoiSetMetaParams(
+                mask_type='countours',
+                filters={'area': {'min': 1e3, 'max': 1e4}},
+                expand_box_by=(128, 2),
+                deproject_channel=0
+            )
+        )
+        roiset.classify_by('permissive_model', [0], IntensityThresholdInstanceMaskSegmentationModel(tr=0.0))
+        self.assertEqual(roiset.get_df()['classify_by_permissive_model'].sum(), roiset.count)
+        roiset.classify_by('avg_intensity', [0], IntensityThresholdInstanceMaskSegmentationModel(tr=0.5))
+        self.assertLess(roiset.get_df()['classify_by_avg_intensity'].sum(), roiset.count)
+        return roiset
+
+    def test_aggregate_classification_results(self):
+        roiset = self.test_roiset_with_instance_segmentation()
+        roiset.aggregate_classifications(
+            query='classify_by_permissive_model == 1 & classify_by_avg_intensity == 1',
+            name='aggregation'
+        )
+        self.assertIn('aggregation', roiset.classification_columns)
+        self.assertTrue(np.all(roiset.get_object_class_map('aggregation').unique()[0] == [0, 1]))
+        self.assertEqual(
+            roiset.get_object_class_map('aggregation').data.sum(),
+            roiset.get_object_class_map('avg_intensity').data.sum(),
+        )
+        self.assertGreater(
+            roiset.get_object_class_map('permissive_model').data.sum(),
+            roiset.get_object_class_map('aggregation').data.sum(),
+        )
+