fimbox/tests/test_generate_dem_diff.py at main · sdmlua/fimbox · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
"""
Tests for bridge DEM processing pipeline.
Step 1 (generateBridgeRaster): streams USGS LiDAR and writes per-bridge .tif
Step 2 (BridgeDEMDiff):        computes lidar_elev - dem_elev and saves bridge_elev_diff.tif
"""

import logging
from pathlib import Path

import fimbox

log = logging.getLogger(__name__)

OUT_DIR = Path(__file__).resolve().parents[2] / "out" / "test_smallB" / "watershed-data"
# OUT_DIR = (
#     Path(__file__).resolve().parents[2]
#     / "out"
#     / "nwm_11239459and2more"
#     / "watershed-data"
# )

bridge_gpkg = OUT_DIR / "osm_bridges_subset.gpkg"  # TODO: change to full bridges gpkg
dem_path = OUT_DIR / "dem.tif"
out_dir = OUT_DIR


# check which bridges already have rasters vs still pending (safe to run anytime)
def test_bridge_raster_status():
    info = fimbox.generateBridgeRaster(
        bridge_gpkg=bridge_gpkg,
        out_dir=out_dir,
    ).status()
    assert "total" in info


# download LiDAR and build per-bridge elevation tifs
def test_generate_bridge_raster():
    tif_dir = fimbox.generateBridgeRaster(
        bridge_gpkg=bridge_gpkg,
        out_dir=out_dir,
        resolution=10.0,
        buffer_m=10.0,
        n_workers=4,
        # id_col="my_id",  # only needed if gpkg has no 'osmid' column
    ).run()
    log.info(f"Per-bridge tifs --> {tif_dir}")


# compute difference raster
def test_bridge_dem_diff():
    out_path = fimbox.BridgeDEMDiff(
        dem_path=dem_path,
        lidar_tif_dir=OUT_DIR / "bridge_dem" / "lidar_osm_rasters",
        bridge_gpkg=bridge_gpkg,
        out_dir=out_dir,
        out_name="bridge_elev_diff.tif",
        n_workers=4,
    ).run()
    log.info(f"Bridge diff raster --> {out_path}")


# Run both steps end-to-end
# def test_full_pipeline():
#     tif_dir = fimbox.generateBridgeRaster(
#         bridge_gpkg=bridge_gpkg,
#         out_dir=out_dir,
#         resolution=10.0,
#         n_workers=4,
#     ).run()
#
#     out_path = fimbox.BridgeDEMDiff(
#         dem_path=dem_path,
#         lidar_tif_dir=tif_dir,
#         bridge_gpkg=bridge_gpkg,
#         out_dir=out_dir,
#         n_workers=4,
#     ).run()
#     print(f"Done: {out_path}")