Add test for predict_aop_data

Author: satsin06

neonwranglerpy/lib/extract_lidar_data.py

@@ -1,4 +1,4 @@
-"""Some text here."""
+"""Function to extract lidar data using RGB data and vst data."""
 import laspy
 import numpy as np
 import os
@@ -13,9 +13,10 @@ def extract_lidar_data(rgb_data,
                        dpID="DP1.30003.001",
                        site="DELA"):
     """
-    Extract LiDAR data with geo_data_frame and image predictions.
+    Extract LiDAR data using RGB data and vst data.
 
     Arguments:
+    ------------
         rgb_data: GeoDataFrame containing the plot data
         vst_data: DataFrame containing the plot data
         year: Year of the data

neonwranglerpy/lib/predict_aop_data.py

@@ -4,6 +4,7 @@
 import os
 import re
 import matplotlib.pyplot as plt
+import pandas as pd
 from shapely.geometry import Point
 import rasterio
 from deepforest import main
@@ -42,8 +43,8 @@ def predict_aop_data(vst_data,
         vst_data['utmZone'].map(lambda x: (326 * 100) + int(x[:-1]))).astype(str)
     geo_data_frame = gpd.GeoDataFrame(vst_data, geometry=geometry, crs=epsg_codes.iloc[0])
     site_level_data = vst_data[vst_data.plotID.str.contains(site)]
-    get_tiles = ((site_level_data.easting / 1000).astype(int) * 1000).astype(str) + "_"
-    +((site_level_data.northing / 1000).astype(int) * 1000).astype(str)
+    get_tiles = (((site_level_data.easting / 1000).astype(int) * 1000).astype(str) + "_" +
+                 ((site_level_data.northing / 1000).astype(int) * 1000).astype(str))
     print(get_tiles.unique())
 
     pattern = fr"{year}_{site}_.*_{get_tiles.unique()[0]}"
@@ -96,4 +97,6 @@ def predict_aop_data(vst_data,
 
                         all_predictions.append(prediction)
 
-    return all_predictions
+    all_predictions_df = pd.concat(all_predictions)
+
+    return all_predictions_df

tests/test_predict_aop_data.py

@@ -0,0 +1,18 @@
+"""Test predict_aop_data.py file."""
+import pandas as pd
+from neonwranglerpy.lib.predict_aop_data import predict_aop_data
+
+
+def test_predict_aop_data():
+    """Test predict_aop_data function."""
+    savepath = 'tests/raw_data'
+    vst_data = pd.read_csv('tests/raw_data/vst_data.csv')
+
+    result = predict_aop_data(vst_data=vst_data.iloc[1:10, :], year='2018',
+                              dpID='DP3.30010.001', savepath=savepath, site='DELA',
+                              plot_crop=False)
+
+    assert (vst_data.shape[0] > 0) & (vst_data.shape[1] > 0)
+    assert len(result) > 0
+    assert isinstance(result, pd.DataFrame)
+    assert result[['xmin', 'ymin', 'xmax', 'ymax']].duplicated().any()