Efficient spatial indexing for clusters using R-tree and interval trees.
| 46 | |
| 47 | |
| 48 | class SpatialClusterIndex: |
| 49 | """Efficient spatial indexing for clusters using R-tree and interval trees.""" |
| 50 | |
| 51 | def __init__(self, clusters: List[Cluster]): |
| 52 | p = index.Property() |
| 53 | p.dimension = 2 |
| 54 | self.spatial_index = index.Index(properties=p) |
| 55 | self.x_intervals = IntervalTree() |
| 56 | self.y_intervals = IntervalTree() |
| 57 | self.clusters_by_id: Dict[int, Cluster] = {} |
| 58 | |
| 59 | for cluster in clusters: |
| 60 | self.add_cluster(cluster) |
| 61 | |
| 62 | def add_cluster(self, cluster: Cluster): |
| 63 | bbox = cluster.bbox |
| 64 | self.spatial_index.insert(cluster.id, bbox.as_tuple()) |
| 65 | self.x_intervals.insert(bbox.l, bbox.r, cluster.id) |
| 66 | self.y_intervals.insert(bbox.t, bbox.b, cluster.id) |
| 67 | self.clusters_by_id[cluster.id] = cluster |
| 68 | |
| 69 | def remove_cluster(self, cluster: Cluster): |
| 70 | self.spatial_index.delete(cluster.id, cluster.bbox.as_tuple()) |
| 71 | del self.clusters_by_id[cluster.id] |
| 72 | |
| 73 | def find_candidates(self, bbox: BoundingBox) -> Set[int]: |
| 74 | """Find potential overlapping cluster IDs using all indexes.""" |
| 75 | spatial = set(self.spatial_index.intersection(bbox.as_tuple())) |
| 76 | x_candidates = self.x_intervals.find_containing( |
| 77 | bbox.l |
| 78 | ) | self.x_intervals.find_containing(bbox.r) |
| 79 | y_candidates = self.y_intervals.find_containing( |
| 80 | bbox.t |
| 81 | ) | self.y_intervals.find_containing(bbox.b) |
| 82 | return spatial.union(x_candidates).union(y_candidates) |
| 83 | |
| 84 | def check_overlap( |
| 85 | self, |
| 86 | bbox1: BoundingBox, |
| 87 | bbox2: BoundingBox, |
| 88 | overlap_threshold: float, |
| 89 | containment_threshold: float, |
| 90 | ) -> bool: |
| 91 | """Check if two bboxes overlap sufficiently.""" |
| 92 | area1, area2 = bbox1.area(), bbox2.area() |
| 93 | if area1 <= 0 or area2 <= 0: |
| 94 | return False |
| 95 | |
| 96 | overlap_area = bbox1.intersection_area_with(bbox2) |
| 97 | if overlap_area <= 0: |
| 98 | return False |
| 99 | |
| 100 | iou = overlap_area / (area1 + area2 - overlap_area) |
| 101 | containment1 = overlap_area / area1 |
| 102 | containment2 = overlap_area / area2 |
| 103 | |
| 104 | return ( |
| 105 | iou > overlap_threshold |