diff --git a/sahi/__init__.py b/sahi/__init__.py index 6d7654ffc..330744e0a 100644 --- a/sahi/__init__.py +++ b/sahi/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.11.16" +__version__ = "0.11.18" from sahi.annotation import BoundingBox, Category, Mask from sahi.auto_model import AutoDetectionModel diff --git a/sahi/utils/coco.py b/sahi/utils/coco.py index e84f2f3d3..942561cb7 100644 --- a/sahi/utils/coco.py +++ b/sahi/utils/coco.py @@ -14,7 +14,7 @@ from typing import Dict, List, Optional, Set, Union import numpy as np -from shapely import MultiPolygon +from shapely import GeometryCollection, MultiPolygon, Polygon from shapely.validation import make_valid from tqdm import tqdm @@ -224,12 +224,27 @@ 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) - if not isinstance(valid, MultiPolygon): - valid = MultiPolygon([valid]) + valid = filter_polygons(valid) self._shapely_annotation.multipolygon = valid intersection_shapely_annotation = self._shapely_annotation.get_intersection(shapely_polygon) return CocoAnnotation.from_shapely_annotation(