drape

Author: bw4sz

drape_CHM.py

@@ -2,22 +2,69 @@
 from src import CHM
 from src.data import read_config
 from src.utils import create_glob_lists
+from src.neon_paths import find_sensor_path
+from src.start_cluster import start
+from src.model_list import species_model_paths
+
 import glob
 import geopandas as gpd
+import rasterstats
+import os
+from distributed import wait, as_completed
+import traceback
 
 config = read_config("config.yml")
 rgb_pool, h5_pool, hsi_pool, CHM_pool = create_glob_lists(config)
 
-files = glob.glob("/blue/ewhite/b.weinstein/DeepTreeAttention/results/predictions/**/*.shp", recursive=True)
+all_files = []
+for site in species_model_paths:
+    model_path = species_model_paths[site]
+    model_name = os.path.splitext(os.path.basename(model_path))[0]
+    prediction_dir = os.path.join("/blue/ewhite/b.weinstein/DeepTreeAttention/results/predictions/{}/{}/*.shp".format(site, model_name))
+    files = glob.glob(prediction_dir, recursive=True)
+    all_files.append(files)
+
+all_files = [item for sublist in all_files for item in sublist]
+print("Found {} files".format(len(all_files)))
+
+client = start(cpus=120,mem_size="4GB")
 
 def drape(shp, config, CHM_pool):
     """Take in a predictions shapefile and extract CHM height"""
-    shp = gpd.read_file(shp)
-    # Create a dummy plotID
-    shp["plotID"] = "same_plot"
-    draped_shp = CHM.postprocess_CHM(shp, lookup_pool=CHM_pool)
-    #draped_shp.to_file(shp)
+    df = gpd.read_file(shp)
+
+    #buffer slightly, CHM model can be patchy
+    geom = df.geometry.buffer(1)
+    try:
+        CHM_path = find_sensor_path(lookup_pool=CHM_pool, bounds=df.total_bounds)
+    except ValueError:
+        return None
+    
+    draped_boxes = rasterstats.zonal_stats(geom,
+                                            CHM_path,
+                                           add_stats={'q99': CHM.non_zero_99_quantile})
+    df["height"] = [x["q99"] for x in draped_boxes]
+    filtered_trees = df[df["height"]>3]
 
+    dirname = os.path.dirname(shp)
+    basename = os.path.basename(shp)
+    dst = os.path.join(dirname,"draped")
+    os.makedirs(dst, exist_ok=True)
+    filtered_trees.to_file(os.path.join(dst, basename))
+
+    return dst
+
+futures = []
 for f in files:
-    drape(f, config=config, CHM_pool=CHM_pool)
-    print(f)
+    future = client.submit(drape, f, CHM_pool=CHM_pool, config=config)
+    futures.append(future)
+
+for x in as_completed(futures):
+    try:
+        print(x.result())
+    except:
+        traceback.print_exc()
+
+#for f in files:
+#    dst = drape(f, config=config, CHM_pool=CHM_pool)
+#    print(dst)

predict.py

@@ -128,6 +128,7 @@ def create_landscape_map(site, model_path, config, client, rgb_pool, hsi_pool, h
                 config=config,
                 dead_model_path=dead_model_path,
                 savedir="/blue/ewhite/b.weinstein/DeepTreeAttention/results/crowns",
+                CHM_pool=CHM_pool,
                 overwrite=False)
         except:
             traceback.print_exc()

src/predict.py

@@ -106,7 +106,7 @@ def generate_prediction_crops(crown_path, config, rgb_pool, h5_pool, img_pool, c
 
     #Write file alongside       
     crown_annotations = gpd.GeoDataFrame(crown_annotations, geometry="geometry")    
-    crown_annotations = crown_annotations.merge(crowns[["individual","dead_label","dead_score", "score"]])
+    crown_annotations = crown_annotations.merge(crowns[["individual","dead_label","dead_score", "score","CHM_height"]])
 
     crown_annotations.to_file(output_name)  
 
@@ -166,8 +166,8 @@ def predict_tile(crown_annotations, m, config, savedir, site, trainer, filter_de
     trees["crown_score"] = trees["score"]
     trees = trees.drop(columns=["pred_taxa_top1","label","score","taxonID"])
     trees = trees.groupby("individual").apply(lambda x: x.head(1)).reset_index(drop=True)
-    trees = gpd.GeoDataFrame(trees, geometry="geometry")    
-    
+    trees = gpd.GeoDataFrame(trees, geometry="geometry")
+
     #Save .shp
     output_name = os.path.splitext(os.path.basename(crown_annotations))[0]
     trees.to_file(os.path.join(savedir, "{}.shp".format(output_name)))