grass_make_canopy_dem.py

#!/usr/bin/env python
############################################################################
#
# MODULE:       grass_make_canopy_dem.py (now WhiteboxTools-based)
# AUTHOR:       Collin Bode, UC Berkeley
#
# PURPOSE:
#   Standalone script to create a canopy DEM from unfiltered XYZ LiDAR tiles.
#   1. Import XYZ files and rasterize using max elevation (canopy surface).
#   2. Merge tiles into a single canopy raster.
#   NOTE: This is a standalone script and does not use ssr_params.py.
#
# COPYRIGHT:    (c) 2015 Collin Bode
#               This program is free software under the GNU General Public
#               License (>=v2). Read the file COPYING that comes with GRASS
#               for details.
#
#############################################################################

import os
import sys
import traceback
import datetime as dt
import numpy as np
import rasterio
from rasterio.transform import Affine
from rasterio.merge import merge as rasterio_merge

# ---------------------------------------------------------------------------
# CONFIGURATION - edit these values for your dataset
# ---------------------------------------------------------------------------
workspace = '/data/ssr_workspace'
C = 1                           # cell size in meters
bregion = 'b8k'
loc = location = 'angelo2014'
P = loc + str(C) + 'm'
pref = bregion + str(C) + 'm'

year = 'y14'
inPath = '/data/source/LiDAR/2014_EelBathymetry_LiDAR/Angelo/Tiles_ASCII_xyz/unfiltered/'
inSuffix = 'xyz'
sep = ','
overlap = float(0.00)
pmaxpref = 'pmax_c' + str(C) + year

dem_name = P + 'dem'
can_name = P + 'can'
vegheight_name = P + 'vegheight'
output_dir = os.path.join(workspace, 'canopy')
# ---------------------------------------------------------------------------


def printout(str_text, lf):
    timestamp = dt.datetime.strftime(dt.datetime.now(), "%H:%M:%S")
    lf.write(timestamp + ": " + str_text + '\n')
    print(timestamp + ": " + str_text)


def ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def xyz_to_raster_max(xyz_path, cell_size_m, overlap_m=0.0, separator=','):
    """Rasterize XYZ point cloud using maximum elevation per cell."""
    data = np.loadtxt(xyz_path, delimiter=separator)
    if data.ndim == 1:
        data = data.reshape(1, -1)
    x, y, z = data[:, 0], data[:, 1], data[:, 2]

    x_min = x.min() - overlap_m
    x_max = x.max() + overlap_m
    y_min = y.min() - overlap_m
    y_max = y.max() + overlap_m

    cols = max(1, int(np.ceil((x_max - x_min) / cell_size_m)))
    rows = max(1, int(np.ceil((y_max - y_min) / cell_size_m)))

    col_idx = np.floor((x - x_min) / cell_size_m).astype(int)
    row_idx = np.floor((y_max - y) / cell_size_m).astype(int)

    valid = (col_idx >= 0) & (col_idx < cols) & (row_idx >= 0) & (row_idx < rows)
    col_idx = col_idx[valid]
    row_idx = row_idx[valid]
    z_valid = z[valid]

    grid = np.full((rows, cols), np.nan, dtype=np.float32)
    for r, c, zv in zip(row_idx, col_idx, z_valid):
        if np.isnan(grid[r, c]) or zv > grid[r, c]:
            grid[r, c] = zv

    transform = Affine(cell_size_m, 0, x_min, 0, -cell_size_m, y_max)
    return grid, transform


def save_raster_tif(data, transform, path, nodata=np.nan):
    """Save array as GeoTIFF with LZW compression."""
    ensure_dir(os.path.dirname(path))
    profile = {
        'driver': 'GTiff',
        'compress': 'lzw',
        'dtype': 'float32',
        'nodata': nodata,
        'width': data.shape[1],
        'height': data.shape[0],
        'count': 1,
        'transform': transform,
    }
    with rasterio.open(path, 'w', **profile) as dst:
        dst.write(data, 1)


def mosaic_tifs(input_paths, output_path):
    """Merge list of GeoTIFF files into one."""
    datasets = [rasterio.open(p) for p in input_paths]
    mosaic, transform = rasterio_merge(datasets)
    profile = datasets[0].profile.copy()
    profile.update(
        driver='GTiff', compress='lzw',
        height=mosaic.shape[1], width=mosaic.shape[2],
        transform=transform, count=1,
    )
    for ds in datasets:
        ds.close()
    ensure_dir(os.path.dirname(output_path))
    with rasterio.open(output_path, 'w', **profile) as dst:
        dst.write(mosaic[0], 1)


def main():
    ensure_dir(output_dir)

    tlog = dt.datetime.strftime(dt.datetime.now(), "%Y-%m-%d_h%H")
    log_path = os.path.join(output_dir, 'canopy_' + tlog + '.log')
    lf = open(log_path, 'a')

    printout("STARTING CANOPY DEM CREATION", lf)
    printout("Location: " + loc, lf)
    printout("LiDAR year: " + year, lf)
    printout("Cell size: " + str(C) + " m", lf)
    printout("Input path: " + inPath, lf)
    printout("Overlap: " + str(overlap) + " m", lf)
    printout('_________________________________', lf)

    L2start = dt.datetime.now()
    printout('START LiDAR import', lf)

    tile_paths = []

    for strFile in sorted(os.listdir(inPath)):
        parts = strFile.split('.')
        if len(parts) < 2 or parts[1] != inSuffix:
            continue

        tile = parts[0]
        xyz_full = os.path.join(inPath, strFile)
        printout("Importing: " + strFile, lf)

        grid, transform = xyz_to_raster_max(xyz_full, float(C), overlap, sep)

        tile_tif = os.path.join(output_dir, 'tile_' + tile + '.tif')
        save_raster_tif(grid, transform, tile_tif)
        tile_paths.append(tile_tif)
        printout("Tile saved: " + tile, lf)

    if not tile_paths:
        printout("ERROR: No tiles found in " + inPath, lf)
        lf.close()
        sys.exit(1)

    printout("Mosaicking " + str(len(tile_paths)) + " tiles", lf)
    output_raster = os.path.join(output_dir, pmaxpref + pref + '.tif')
    mosaic_tifs(tile_paths, output_raster)

    for tp in tile_paths:
        os.remove(tp)

    L2end = dt.datetime.now()
    printout("DONE. Output: " + output_raster, lf)
    printout("Processing time: " + str(L2end - L2start), lf)
    printout("--------------------------------------", lf)
    lf.close()
    sys.exit("FINISHED.")


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()
        sys.exit(1)