From e57bd7e092a4c57a7d969680b1926560d40e6eae Mon Sep 17 00:00:00 2001
From: Richard <richard@nieding.de>
Date: Thu, 14 Dec 2023 16:08:24 +0100
Subject: [PATCH] added additional cleaning functions for very small polygons

---
 detectree2/models/outputs.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/detectree2/models/outputs.py b/detectree2/models/outputs.py
index 679c0e39..a057e385 100644
--- a/detectree2/models/outputs.py
+++ b/detectree2/models/outputs.py
@@ -458,6 +458,26 @@ def post_clean(unclean_df: gpd.GeoDataFrame,
 
     return reclean_df.reset_index(drop=True)
 
+def remove_very_small_polygons(crowns:gpd.GeoDataFrame, size_threshold=1.0) -> gpd.GeoDataFrame:
+    """Removes small fragmentary crowns below or equal to the threshold.
+
+    Args:
+        crowns (gpd.GeoDataFrame): A GeoDataFrame with small fragmentary crowns.
+        size_threshold (float, optional): given in m² (square meter). Defaults to 1.0.
+
+    Returns:
+        gpd.GeoDataFrame: A GeoDataFrame with no small fragmentary crowns.
+    """
+    areas = [(row.geometry.area, index) for index, row in crowns.iterrows()]
+    areas.sort()
+    indexes_to_remove = [area[1] for area in areas if area[0] <= size_threshold] # threshold may differ across inference images
+    indexes_to_remove = list(set(indexes_to_remove))
+    indexes_to_keep = set(range(crowns.shape[0])) - set(indexes_to_remove)
+    print("Removed",len(indexes_to_remove),"very small crowns!")
+    output_gdf = crowns.take(list(indexes_to_keep))
+    output_gdf = output_gdf.reset_index(drop=True)
+    return output_gdf
+
 
 def load_geopandas_dataframes(folder):
     """Load all GeoPackage files in a folder into a list of GeoDataFrames."""