Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix TypeError: 'GeometryCollection' object is not subscriptable when slicing COCO #1047

Merged
merged 4 commits into from
Jun 2, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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