fix(sahi): Fix Polygon Repair and Empty Polygon Issues

sahi/postprocess/utils.py

@@ -4,6 +4,7 @@
 import numpy as np
 import torch
 from shapely.geometry import MultiPolygon, Polygon
+from shapely.geometry.collection import GeometryCollection
 
 from sahi.annotation import BoundingBox, Category, Mask
 from sahi.prediction import ObjectPrediction
@@ -64,6 +65,81 @@ def tolist(self):
         else:
             return self.list
 
+###################################################sunhao###################################################
+def repair_polygon(
+    shapely_polygon: Polygon
+) -> Polygon:
+    """
+    Fix polygons
+
+    :param shapely_polygon: Shapely polygon object
+    :return:
+    """
+    if not shapely_polygon.is_valid:
+        fixed_polygon = shapely_polygon.buffer(0)
+        if fixed_polygon.is_valid:
+            if isinstance(fixed_polygon, Polygon):
+                return fixed_polygon
+            elif isinstance(fixed_polygon, MultiPolygon):
+                return max(fixed_polygon.geoms, key=lambda p: p.area)
+            elif isinstance(fixed_polygon, GeometryCollection):
+                polygons = [geom for geom in fixed_polygon.geoms if isinstance(geom, Polygon)]
+                return max(polygons, key=lambda p: p.area) if polygons else shapely_polygon
+
+    return shapely_polygon
+
+
+def repair_multipolygon(
+    shapely_multipolygon: MultiPolygon
+) -> MultiPolygon:
+    """
+    Fix invalid MultiPolygon objects
+
+    :param shapely_multipolygon: Imported shapely MultiPolygon object
+    :return:
+    """
+    if not shapely_multipolygon.is_valid:
+        fixed_geometry = shapely_multipolygon.buffer(0)
+
+        if fixed_geometry.is_valid:
+            if isinstance(fixed_geometry, MultiPolygon):
+                return fixed_geometry
+            elif isinstance(fixed_geometry, Polygon):
+                return MultiPolygon([fixed_geometry])
+            elif isinstance(fixed_geometry, GeometryCollection):
+                polygons = [geom for geom in fixed_geometry.geoms if isinstance(geom, Polygon)]
+                return MultiPolygon(polygons) if polygons else shapely_multipolygon
+
+    return shapely_multipolygon
+
+
+def coco_segmentation_to_shapely(
+    segmentation: Union[list, list[list]]
+):
+    """
+    Fix segment data in COCO format
+
+    :param segmentation: segment data in COCO format
+    :return:
+    """
+    if isinstance(segmentation, list) and all([not isinstance(seg, list) for seg in segmentation]):
+        segmentation = [segmentation]
+    elif isinstance(segmentation, list) and all([isinstance(seg, list) for seg in segmentation]):
+        pass
+    else:
+        raise ValueError("segmentation must be list or list[list]")
+
+    polygon_list = []
+
+    for coco_polygon in segmentation:
+        point_list = list(zip(coco_polygon[::2], coco_polygon[1::2]))
+        shapely_polygon = Polygon(point_list)
+        polygon_list.append(repair_polygon(shapely_polygon))
+
+    shapely_multipolygon = repair_multipolygon(MultiPolygon(polygon_list))
+    return shapely_multipolygon
+
+############################################################################################################
 
 def object_prediction_list_to_torch(object_prediction_list: ObjectPredictionList) -> torch.tensor:
     """
@@ -166,6 +242,12 @@ def get_merged_mask(pred1: ObjectPrediction, pred2: ObjectPrediction) -> Mask:
     # buffer(0) is a quickhack to fix invalid polygons most of the time
     poly1 = get_shapely_multipolygon(mask1.segmentation).buffer(0)
     poly2 = get_shapely_multipolygon(mask2.segmentation).buffer(0)
+    ###################################################sunhao###################################################
+    if poly1.is_empty:
+        poly1 = coco_segmentation_to_shapely(mask1.segmentation)
+    if poly2.is_empty:
+        poly2 = coco_segmentation_to_shapely(mask2.segmentation)
+    ############################################################################################################
     union_poly = poly1.union(poly2)
     if not hasattr(union_poly, "geoms"):
         union_poly = MultiPolygon([union_poly])

sahi/slicing.py

@@ -35,6 +35,7 @@ def get_slice_bboxes(
     auto_slice_resolution: bool = True,
     overlap_height_ratio: float = 0.2,
     overlap_width_ratio: float = 0.2,
+    inference_org_image: bool = False
 ) -> List[List[int]]:
     """Slices `image_pil` in crops.
     Corner values of each slice will be generated using the `slice_height`,
@@ -53,6 +54,7 @@ def get_slice_bboxes(
             overlap of 20 pixels). Default 0.2.
         auto_slice_resolution (bool): if not set slice parameters such as slice_height and slice_width,
             it enables automatically calculate these params from image resolution and orientation.
+        inference_org_image (bool): 如果为True, 则在裁剪列表的最后将原图加入
 
     Returns:
         List[List[int]]: List of 4 corner coordinates for each N slices.
@@ -88,7 +90,11 @@ def get_slice_bboxes(
                 slice_bboxes.append([x_min, y_min, x_max, y_max])
             x_min = x_max - x_overlap
         y_min = y_max - y_overlap
-    return slice_bboxes
+
+    if inference_org_image:
+        return slice_bboxes + [[0, 0, image_width, image_height]]
+    else:
+        return slice_bboxes
 
 
 def annotation_inside_slice(annotation: Dict, slice_bbox: List[int]) -> bool:
@@ -236,10 +242,10 @@ def filenames(self) -> List[int]:
     def __getitem__(self, i):
         def _prepare_ith_dict(i):
             return {
-                "image": self.images[i],
-                "coco_image": self.coco_images[i],
+                "image"         : self.images[i],
+                "coco_image"    : self.coco_images[i],
                 "starting_pixel": self.starting_pixels[i],
-                "filename": self.filenames[i],
+                "filename"      : self.filenames[i],
             }
 
         if isinstance(i, np.ndarray):
@@ -273,6 +279,7 @@ def slice_image(
     min_area_ratio: float = 0.1,
     out_ext: Optional[str] = None,
     verbose: bool = False,
+    inference_org_image: bool = False
 ) -> SliceImageResult:
     """Slice a large image into smaller windows. If output_file_name is given export
     sliced images.
@@ -299,6 +306,7 @@ def slice_image(
             original suffix for lossless image formats and png for lossy formats ('.jpg','.jpeg').
         verbose (bool, optional): Switch to print relevant values to screen.
             Default 'False'.
+        inference_org_image (bool): If True, the original image is used for inference. Default False.
 
     Returns:
         sliced_image_result: SliceImageResult:
@@ -341,6 +349,7 @@ def _export_single_slice(image: np.ndarray, output_dir: str, slice_file_name: st
         slice_width=slice_width,
         overlap_height_ratio=overlap_height_ratio,
         overlap_width_ratio=overlap_width_ratio,
+        inference_org_image=inference_org_image
     )
 
     n_ims = 0