diff --git a/sahi/models/detectron2.py b/sahi/models/detectron2.py index 223a6d825..e8ce35d98 100644 --- a/sahi/models/detectron2.py +++ b/sahi/models/detectron2.py @@ -149,9 +149,9 @@ def _create_object_prediction_list_from_original_predictions( object_prediction_list = [ ObjectPrediction( bbox=box.tolist() if mask is None else None, - segmentation=get_coco_segmentation_from_bool_mask(mask.detach().cpu().numpy()) - if mask is not None - else None, + segmentation=( + get_coco_segmentation_from_bool_mask(mask.detach().cpu().numpy()) if mask is not None else None + ), category_id=category_id.item(), category_name=self.category_mapping[str(category_id.item())], shift_amount=shift_amount, diff --git a/sahi/utils/coco.py b/sahi/utils/coco.py index 942561cb7..cfc5c6de5 100644 --- a/sahi/utils/coco.py +++ b/sahi/utils/coco.py @@ -224,28 +224,7 @@ def __init__( self._shapely_annotation = shapely_annotation def get_sliced_coco_annotation(self, slice_bbox: List[int]): - def filter_polygons(geometry): - """ - Filters out and returns only Polygon or MultiPolygon components of a geometry. - If geometry is a Polygon, it converts it into a MultiPolygon. - If it's a GeometryCollection, it filters to create a MultiPolygon from any Polygons in the collection. - Returns an empty MultiPolygon if no Polygon or MultiPolygon components are found. - """ - if isinstance(geometry, Polygon): - return MultiPolygon([geometry]) - elif isinstance(geometry, MultiPolygon): - return geometry - elif isinstance(geometry, GeometryCollection): - polygons = [geom for geom in geometry.geoms if isinstance(geom, Polygon)] - return MultiPolygon(polygons) if polygons else MultiPolygon() - return MultiPolygon() - shapely_polygon = box(slice_bbox[0], slice_bbox[1], slice_bbox[2], slice_bbox[3]) - samp = self._shapely_annotation.multipolygon - if not samp.is_valid: - valid = make_valid(samp) - valid = filter_polygons(valid) - self._shapely_annotation.multipolygon = valid intersection_shapely_annotation = self._shapely_annotation.get_intersection(shapely_polygon) return CocoAnnotation.from_shapely_annotation( intersection_shapely_annotation, diff --git a/sahi/utils/shapely.py b/sahi/utils/shapely.py index e579a0851..7380e9f7b 100644 --- a/sahi/utils/shapely.py +++ b/sahi/utils/shapely.py @@ -3,7 +3,8 @@ from typing import List -from shapely.geometry import CAP_STYLE, JOIN_STYLE, MultiPolygon, Polygon, box +from shapely.geometry import CAP_STYLE, JOIN_STYLE, GeometryCollection, MultiPolygon, Polygon, box +from shapely.validation import make_valid def get_shapely_box(x: int, y: int, width: int, height: int) -> Polygon: @@ -21,8 +22,26 @@ def get_shapely_box(x: int, y: int, width: int, height: int) -> Polygon: def get_shapely_multipolygon(coco_segmentation: List[List]) -> MultiPolygon: """ - Accepts coco style polygon coords and converts it to shapely multipolygon object + Accepts coco style polygon coords and converts it to valid shapely multipolygon object """ + + def filter_polygons(geometry): + """ + Filters out and returns only Polygon or MultiPolygon components of a geometry. + If geometry is a Polygon, it converts it into a MultiPolygon. + If it's a GeometryCollection, it filters + to create a MultiPolygon from any Polygons in the collection. + Returns an empty MultiPolygon if no Polygon or MultiPolygon components are found. + """ + if isinstance(geometry, Polygon): + return MultiPolygon([geometry]) + elif isinstance(geometry, MultiPolygon): + return geometry + elif isinstance(geometry, GeometryCollection): + polygons = [geom for geom in geometry.geoms if isinstance(geom, Polygon)] + return MultiPolygon(polygons) if polygons else MultiPolygon() + return MultiPolygon() + polygon_list = [] for coco_polygon in coco_segmentation: point_list = list(zip(coco_polygon[0::2], coco_polygon[1::2])) @@ -30,6 +49,9 @@ def get_shapely_multipolygon(coco_segmentation: List[List]) -> MultiPolygon: polygon_list.append(shapely_polygon) shapely_multipolygon = MultiPolygon(polygon_list) + if not shapely_multipolygon.is_valid: + shapely_multipolygon = filter_polygons(make_valid(shapely_multipolygon)) + return shapely_multipolygon diff --git a/tests/test_shapelyutils.py b/tests/test_shapelyutils.py index 972402609..d3fa70a95 100644 --- a/tests/test_shapelyutils.py +++ b/tests/test_shapelyutils.py @@ -26,6 +26,12 @@ def test_get_shapely_multipolygon(self): self.assertEqual(shapely_multipolygon.area, 41177.5) self.assertTupleEqual(shapely_multipolygon.bounds, (1, 1, 325, 200)) + def test_get_shapely_multipolygon_naughty(self): + # self-intersection case + coco_segmentation = [[3559.0, 2046.86, 3.49, 2060.0, 3540.9, 3249.7, 2060.0, 3239.61, 2052.87]] + shapely_multipolygon = get_shapely_multipolygon(coco_segmentation) + self.assertTrue(shapely_multipolygon.is_valid) + def test_shapely_annotation(self): # init shapely_annotation from coco segmentation segmentation = [[1, 1, 325, 125.2, 250, 200, 5, 200]]