feat: add simplify param to merging

Author: okunator

cellseg_models_pytorch/wsi/inst_merger.py

@@ -8,8 +8,6 @@
 from shapely.geometry import LineString, Polygon, box
 from tqdm import tqdm
 
-from cellseg_models_pytorch.utils.spatial_ops import get_objs
-
 __all__ = ["InstMerger"]
 
 
@@ -30,12 +28,17 @@ def __init__(
         self.grid = gpd.GeoDataFrame({"geometry": polygons})
         self.gdf = gdf
 
-    def merge(self, dst: str = None) -> Union[gpd.GeoDataFrame, None]:
+    def merge(
+        self, dst: str = None, simplify_level: int = 1
+    ) -> Union[gpd.GeoDataFrame, None]:
         """Merge the instances at the image boundaries.
 
         Parameters:
             dst (str):
                 The destination directory to save the merged instances.
+                If None, the merged GeoDataFrame is returned.
+            simplify_level (int, default=1):
+                The level of simplification to apply to the merged instances.
 
         Returns:
             Union[gpd.GeoDataFrame, None]:
@@ -63,6 +66,7 @@ def merge(self, dst: str = None) -> Union[gpd.GeoDataFrame, None]:
         merged = pd.concat([merge_obj_x, merge_obj_y, non_boundary_objs]).reset_index(
             drop=True
         )
+        merged.geometry = merged.geometry.simplify(simplify_level)
 
         if dst is not None:
             if suff == ".parquet":
@@ -117,6 +121,21 @@ def _merge_boundary_objs(
 
         return merged.reset_index(drop=True)
 
+    def _get_objs(
+        self,
+        objects: gpd.GeoDataFrame,
+        area: gpd.GeoDataFrame,
+        predicate: str,
+        **kwargs,
+    ) -> gpd.GeoDataFrame:
+        """Get the objects that intersect with the midline."""
+        inds = objects.geometry.sindex.query(
+            area.geometry, predicate=predicate, **kwargs
+        )
+        objs: gpd.GeoDataFrame = objects.iloc[np.unique(inds)[1:]]
+
+        return objs.drop_duplicates("geometry")
+
     def _merge_objs_axis(
         self,
         grid: gpd.GeoDataFrame,
@@ -148,7 +167,7 @@ def _merge_objs_axis(
         desc = "Merging objects (x-axis)" if axis == "x" else "Merging objects (y-axis)"
         for start, next_col in tqdm(grouped_list[1:], desc=desc):
             grid_union = pd.concat([last_col, next_col]).union_all()
-            objs = get_objs(self._union_to_gdf(grid_union), gdf, predicate)
+            objs = self._get_objs(gdf, self._union_to_gdf(grid_union), predicate)
 
             minx, miny, maxx, maxy = next_col.total_bounds
 
@@ -164,12 +183,12 @@ def _merge_objs_axis(
                 midline_gdf = gpd.GeoDataFrame(geometry=[midline])
 
             # get the cells hitting the midline
-            boundary_objs = get_objs(midline_gdf, objs, predicate)
+            boundary_objs = self._get_objs(objs, midline_gdf, predicate)
 
             non_boundary_objs_left = None
             if get_non_boundary_objs:
-                non_boundary_objs_left = get_objs(
-                    last_col.buffer(-midline_buffer), objs, "contains"
+                non_boundary_objs_left = self._get_objs(
+                    objs, last_col.buffer(-midline_buffer), "contains"
                 )
 
             # merge the boundary objects
@@ -211,15 +230,11 @@ def _get_classes(
         class_names = []
         for ix, row in merged.iterrows():
             area = gpd.GeoDataFrame(geometry=[row.geometry])
-            objs = get_objs(area, non_merged, predicate="intersects")
+            objs = self._get_objs(non_merged, area, predicate="intersects")
 
             if objs.empty:
                 continue
 
-            # HACK: for some reason the first object is always the same
-            if len(objs) > 1 and ix != 0:
-                objs = objs.drop(index=0)
-
             class_names.append(objs.loc[objs.area.idxmax()]["class_name"])
 
         return class_names