Add containment check to clean_crowns, remove remove_contained_crowns

Restores the containment check inside clean_crowns where it belongs.
When a small crown is mostly inside a larger one (default 85%), the
pair is merged and the most confident crown wins — exactly like IoU.
This means:
- Small confident crown inside large less-confident one → large one
  removed → post_clean fills the gap with other crowns
- Small less-confident crown inside large confident one → small
  nested duplicate removed

Removes the separate remove_contained_crowns function which wrongly
always kept the larger crown regardless of confidence.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: PatBall1

detectree2/models/outputs.py

@@ -325,12 +325,20 @@ def clean_crowns(
     confidence=0.2,
     area_threshold=2,
     field="Confidence_score",
+    containment_threshold=0.85,
     verbose=True,
 ) -> gpd.GeoDataFrame:
     """
     Clean overlapping crowns by first identifying all candidate overlapping pairs via a spatial join,
     then clustering crowns into connected components (where an edge is added if two crowns have IoU
-    above a threshold), and finally keeping the best crown (by confidence or any given field) in each cluster.
+    above a threshold, or one crown is largely contained within the other), and finally keeping the
+    best crown (by confidence or any given field) in each cluster.
+
+    The containment check is important because IoU alone misses nested crowns: when a small crown
+    sits inside a much larger one, IoU is low (the union is dominated by the large crown). The
+    containment check merges these pairs so that the most confident crown wins. If the small crown
+    is more confident, the large one is removed — and post_clean can then fill the gap with other
+    crowns.
 
     Args:
         crowns (gpd.GeoDataFrame): Crowns to be cleaned.
@@ -340,6 +348,9 @@ def clean_crowns(
         area_threshold (float, optional): Minimum area of crowns to be retained. Defaults to 2m2 (assuming UTM).
         field (str): Field to used to prioritise selection of crowns. Defaults to "Confidence_score" but this should
             be changed to "Area" if using a model that outputs area.
+        containment_threshold (float, optional): Threshold for the containment ratio
+            (intersection area / smaller crown area). When exceeded, the pair is merged and the most
+            confident crown wins — just like IoU. Set to None to disable. Defaults to 0.85.
 
     Returns:
         gpd.GeoDataFrame: Cleaned crowns.
@@ -376,7 +387,8 @@ def union(x, y):
         if rx != ry:
             parent[ry] = rx
 
-    # 4. For each candidate pair, compute IoU and, if it exceeds the threshold, merge the groups.
+    # 4. For each candidate pair, check IoU and containment; merge if either exceeds its threshold.
+    #    The most confident crown in each merged cluster always wins (step 6).
     for idx, row in tqdm(
         join.iterrows(),
         total=len(join),
@@ -389,10 +401,23 @@ def union(x, y):
         # To avoid duplicate work, skip if i and j are already in the same group.
         if find(i) == find(j):
             continue
-        # Compute the IoU for the pair.
-        iou_val = calc_iou(crowns.at[i, "geometry"], crowns.at[j, "geometry"])
+
+        geom_i = crowns.at[i, "geometry"]
+        geom_j = crowns.at[j, "geometry"]
+        intersection_area = geom_i.intersection(geom_j).area
+
+        # IoU check
+        union_area = geom_i.area + geom_j.area - intersection_area
+        iou_val = intersection_area / union_area if union_area > 0 else 0
         if iou_val > iou_threshold:
             union(i, j)
+            continue
+
+        # Containment check: is the smaller crown mostly inside the larger one?
+        if containment_threshold is not None:
+            min_area = min(geom_i.area, geom_j.area)
+            if min_area > 0 and (intersection_area / min_area) > containment_threshold:
+                union(i, j)
 
     # 5. Group crowns by their union-find root.
     groups: Dict[int, List] = {}
@@ -481,78 +506,6 @@ def post_clean(unclean_df: gpd.GeoDataFrame,
     return current_clean
 
 
-def remove_contained_crowns(
-    crowns: gpd.GeoDataFrame,
-    containment_threshold: float = 0.85,
-    verbose: bool = True,
-) -> gpd.GeoDataFrame:
-    """Remove small crowns that are largely contained within larger crowns.
-
-    This is intended as a final cleanup step *after* clean_crowns and post_clean have run.
-    It catches nested duplicates that IoU-based cleaning misses: when a small crown sits inside
-    a much larger one, IoU is low (the union is dominated by the large crown) but the small
-    crown adds visual clutter rather than useful information.
-
-    The larger crown is always kept. Only the smaller (contained) crown is removed.
-
-    Args:
-        crowns (gpd.GeoDataFrame): Cleaned crowns to check for nesting.
-        containment_threshold (float, optional): Proportion of the smaller crown's area that must
-            fall inside the larger crown for it to be considered contained. Defaults to 0.85.
-        verbose (bool, optional): Print progress information. Defaults to True.
-
-    Returns:
-        gpd.GeoDataFrame: Crowns with nested duplicates removed.
-    """
-    crowns = crowns.copy()
-    crowns["geometry"] = crowns.geometry.buffer(0)
-    crowns.reset_index(drop=True, inplace=True)
-
-    # Spatial join to find candidate overlapping pairs
-    join = gpd.sjoin(crowns, crowns, how="inner", predicate="intersects")
-    join = join[join.index != join.index_right]
-
-    to_remove = set()
-    for _, row in tqdm(
-        join.iterrows(),
-        total=len(join),
-        desc="remove_contained_crowns: checking pairs",
-        smoothing=0,
-        disable=not verbose,
-    ):
-        i = row.name
-        j = row["index_right"]
-
-        # Skip if either crown is already marked for removal
-        if i in to_remove or j in to_remove:
-            continue
-
-        geom_i = crowns.at[i, "geometry"]
-        geom_j = crowns.at[j, "geometry"]
-
-        # Determine which is smaller
-        if geom_i.area <= geom_j.area:
-            small_idx, small_geom, large_geom = i, geom_i, geom_j
-        else:
-            small_idx, small_geom, large_geom = j, geom_j, geom_i
-
-        if small_geom.area == 0:
-            to_remove.add(small_idx)
-            continue
-
-        intersection_area = small_geom.intersection(large_geom).area
-        containment_ratio = intersection_area / small_geom.area
-
-        if containment_ratio > containment_threshold:
-            to_remove.add(small_idx)
-
-    if verbose:
-        print(f"remove_contained_crowns: removed {len(to_remove)} nested crowns "
-              f"({len(crowns)} → {len(crowns) - len(to_remove)})")
-
-    return crowns.drop(index=to_remove).reset_index(drop=True)
-
-
 def load_geopandas_dataframes(folder):
     """Load all GeoPackage files in a folder into a list of GeoDataFrames."""
     all_files = glob.glob(f"{folder}/*.gpkg")