diff --git a/model_server/base/roiset.py b/model_server/base/roiset.py index 998b9e635dcb6dbfd2fd8892f98afe053e7d1871..e1b1c7f397120a2b53508f941e3ef4e07605079d 100644 --- a/model_server/base/roiset.py +++ b/model_server/base/roiset.py @@ -128,6 +128,8 @@ def filter_df_overlap_bbox(df1: pd.DataFrame, df2: pd.DataFrame = None) -> pd.Da :param df1: DataFrame with potentially overlapping bounding boxes :param df2: (optional) second DataFrame :return DataFrame describing subset of overlapping ROIs + bbox_overlaps_with: index of ROI that overlaps + bbox_intersec: pixel area of intersecting region """ def _compare(r0, r1): @@ -136,14 +138,21 @@ def filter_df_overlap_bbox(df1: pd.DataFrame, df2: pd.DataFrame = None) -> pd.Da olz = (r0.zi == r1.zi) return olx and oly and olz + def _intersec(r0, r1): + return (r0.x1 - r1.x0) * (r0.y1 - r1.y0) + first = [] second = [] + intersec = [] if df2 is not None: for pair in itertools.product(df1.index, df2.index): if _compare(df1.iloc[pair[0]], df2.iloc[pair[1]]): first.append(pair[0]) second.append(pair[1]) + intersec.append( + _intersec(df1.iloc[pair[0]], df2.iloc[pair[1]]) + ) else: for pair in itertools.combinations(df1.index, 2): if _compare(df1.iloc[pair[0]], df1.iloc[pair[1]]): @@ -151,9 +160,13 @@ def filter_df_overlap_bbox(df1: pd.DataFrame, df2: pd.DataFrame = None) -> pd.Da second.append(pair[1]) first.append(pair[1]) second.append(pair[0]) + isc = _intersec(df1.iloc[pair[0]], df1.iloc[pair[1]]) + intersec.append(isc) + intersec.append(isc) sdf = df1.iloc[first] sdf.loc[:, 'bbox_overlaps_with'] = second + sdf.loc[:, 'bbox_intersec'] = intersec return sdf # TODO: option to quantify overlap e.g. IOU diff --git a/tests/base/test_roiset.py b/tests/base/test_roiset.py index 91eda9bc171622a15e889d825bb0f66257103c2a..80da0bc9438062cd8f9d48fe3f60759bce8684f8 100644 --- a/tests/base/test_roiset.py +++ b/tests/base/test_roiset.py @@ -644,6 +644,7 @@ class TestRoiSetPolygons(BaseTestRoiSetMonoProducts, unittest.TestCase): self.assertEqual(len(res), 4) self.assertTrue((res.loc[0, 'bbox_overlaps_with'] == [1]).all()) self.assertTrue((res.loc[1, 'bbox_overlaps_with'] == [0, 2]).all()) + self.assertTrue((res.bbox_intersec == 2).all()) return res # TODO: test overloaded condition comparing two DFs @@ -665,6 +666,7 @@ class TestRoiSetPolygons(BaseTestRoiSetMonoProducts, unittest.TestCase): res = filter_df_overlap_bbox(df1, df2) self.assertTrue((res.loc[1, 'bbox_overlaps_with'] == [0]).all()) self.assertEqual(len(res), 1) + self.assertTrue((res.bbox_intersec == 2).all()) def test_overlap_seg(self): df = pd.DataFrame({