ssr_lidar.py

#!/usr/bin/env python
############################################################################
#
# MODULE:       ssr_lidar.py
# AUTHOR:       Collin Bode, UC Berkeley
#
# PURPOSE:
#   1. Accept ASCII xyz LiDAR files (filtered & unfiltered) and rasterize them.
#      Tiles sometimes overlap, so an overlap distance clips tiles before merging.
#   2. Calculate point density rasters.
#   3. Calculate canopy raster using point maximum elevation.
#
# COPYRIGHT:    (c) 2011 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 ssr_params import *
from ssr_utilities import (printout, ensure_dir, raster_path, raster_exists,
                            read_raster, save_raster, mosaic_rasters, setup_workspace)


def xyz_to_raster(xyz_path, cell_size_m, method='count', overlap_m=0.0, separator=','):
    """Bin an XYZ ASCII point cloud into a raster grid.

    method: 'count' for point density, 'max' for canopy surface elevation.
    Returns (data array float32, Affine transform).
    """
    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.zeros((rows, cols), dtype=np.float32)

    if method == 'count':
        np.add.at(grid, (row_idx, col_idx), 1)
    elif method == 'max':
        grid[:] = np.nan
        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 clip_raster_to_extent(data, transform, clip_x_min, clip_x_max, clip_y_min, clip_y_max):
    """Crop a raster array to a tighter bounding box (removes overlap)."""
    cell = transform.a  # cell size
    orig_x_min = transform.c
    orig_y_max = transform.f

    col_start = max(0, int(np.ceil((clip_x_min - orig_x_min) / cell)))
    col_end = min(data.shape[1], int(np.floor((clip_x_max - orig_x_min) / cell)))
    row_start = max(0, int(np.ceil((orig_y_max - clip_y_max) / cell)))
    row_end = min(data.shape[0], int(np.floor((orig_y_max - clip_y_min) / cell)))

    clipped = data[row_start:row_end, col_start:col_end]
    new_transform = Affine(cell, 0, orig_x_min + col_start * cell,
                           0, -cell, orig_y_max - row_start * cell)
    return clipped, new_transform


def build_raster_profile(transform, shape, crs=None):
    """Build a minimal rasterio profile dict."""
    profile = {
        'driver': 'GTiff',
        'compress': 'lzw',
        'dtype': 'float32',
        'nodata': np.nan,
        'width': shape[1],
        'height': shape[0],
        'count': 1,
        'transform': transform,
    }
    if crs:
        profile['crs'] = crs
    return profile


def main():
    setup_workspace()

    tlog = dt.datetime.strftime(dt.datetime.now(), "%Y-%m-%d_h%H")
    log_path = os.path.join(workspace, dir_logs, 'ssr_' + tlog + '_lidar.log')
    ensure_dir(os.path.dirname(log_path))
    lf = open(log_path, 'a')

    cell_size_m = float(C)
    ow = int(lidar_run - 1)   # 0 = no overwrite, 1 = overwrite

    boocan_global = (cansource == '')

    printout("STARTING LiDAR RUN", lf)
    printout("LiDAR year: " + year, lf)
    printout("Prefix: " + P, lf)
    printout("Point Cloud Path: " + inPath, lf)
    printout("Separator: " + sep, lf)
    printout("Tile Overlap (m): " + str(overlap), lf)
    printout("Create Canopy Raster: " + str(boocan_global), lf)
    printout('_________________________________', lf)

    if lidar_run > 0:
        L2start = dt.datetime.now()
        printout('START LiDAR point cloud import', lf)

        lidar_dir = os.path.join(workspace, dir_lidar)
        ensure_dir(lidar_dir)

        for pref in LidarPoints:
            boocan = boocan_global and (pref == 'unfiltered' or pref == 'all')
            printout(f'Processing "{pref}", create canopy: {boocan}', lf)

            inDir = os.path.join(inPath, pref)
            tile_paths = []
            can_tile_paths = []

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

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

                # Point density tile
                density_data, density_transform = xyz_to_raster(
                    xyz_full, cell_size_m, method='count',
                    overlap_m=overlap, separator=sep)

                if overlap > 0:
                    x_min = density_transform.c + overlap
                    x_max = x_min + density_data.shape[1] * cell_size_m - 2 * overlap
                    y_max = density_transform.f - overlap
                    y_min = y_max - density_data.shape[0] * cell_size_m + 2 * overlap
                    density_data, density_transform = clip_raster_to_extent(
                        density_data, density_transform, x_min, x_max, y_min, y_max)

                tile_tif = os.path.join(lidar_dir, 'tile_' + pref + '_' + tile + '.tif')
                profile = build_raster_profile(density_transform, density_data.shape)
                save_raster(density_data, profile, tile_tif)
                tile_paths.append(tile_tif)

                # Canopy tile (max elevation)
                if boocan:
                    can_data, can_transform = xyz_to_raster(
                        xyz_full, cell_size_m, method='max',
                        overlap_m=overlap, separator=sep)
                    if overlap > 0:
                        x_min = can_transform.c + overlap
                        x_max = x_min + can_data.shape[1] * cell_size_m - 2 * overlap
                        y_max = can_transform.f - overlap
                        y_min = y_max - can_data.shape[0] * cell_size_m + 2 * overlap
                        can_data, can_transform = clip_raster_to_extent(
                            can_data, can_transform, x_min, x_max, y_min, y_max)

                    can_tile_tif = os.path.join(lidar_dir, 'cantile_' + tile + '.tif')
                    profile_can = build_raster_profile(can_transform, can_data.shape)
                    save_raster(can_data, profile_can, can_tile_tif)
                    can_tile_paths.append(can_tile_tif)

            # Mosaic all tiles into point density raster
            pointdensity = pdensitypref + pref
            out_density = raster_path(pointdensity, dir_lidar)
            if ow or not os.path.exists(out_density):
                printout('Mosaicking ' + str(len(tile_paths)) + ' tiles -> ' + out_density, lf)
                mosaic_rasters(tile_paths, out_density)
                printout('Point density raster created: ' + pointdensity, lf)

            # Remove individual tiles
            for tp in tile_paths:
                os.remove(tp)

            # Mosaic canopy tiles
            if boocan and can_tile_paths:
                out_can_tmp = raster_path(can + '_temp', dir_lidar)
                mosaic_rasters(can_tile_paths, out_can_tmp)

                # Fill NoData gaps with DEM values (pools, rivers, etc.)
                tmp_data, tmp_profile = read_raster(out_can_tmp)
                dem_data, _ = read_raster(dem_path)

                # Align dem to canopy grid if shapes differ
                if dem_data.shape != tmp_data.shape:
                    printout("WARNING: DEM and canopy shapes differ; using canopy NoData as-is", lf)
                    can_filled = tmp_data
                else:
                    can_filled = np.where(np.isnan(tmp_data), dem_data, tmp_data)

                out_can = raster_path(can, dir_lidar)
                save_raster(can_filled.astype(np.float32), tmp_profile, out_can)
                os.remove(out_can_tmp)
                printout('Canopy raster created: ' + out_can, lf)

                for cp in can_tile_paths:
                    os.remove(cp)

            printout("Done with " + pref, lf)

        L2end = dt.datetime.now()
        printout("DONE with LiDAR import, 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)