Skip to content

Commit

Permalink
Fix TypeError: 'GeometryCollection' object is not subscriptable whe…
Browse files Browse the repository at this point in the history
…n slicing COCO (#1047)
  • Loading branch information
Alias-z authored Jun 2, 2024
1 parent 9503d1b commit 8c07a06
Showing 1 changed file with 18 additions and 3 deletions.
21 changes: 18 additions & 3 deletions sahi/utils/coco.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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(
Expand Down

0 comments on commit 8c07a06

Please sign in to comment.