Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
235 lines (186 sloc) 7.84 KB
"""
Raster data handlers
"""
import logging
import psycopg2 as pg
import numpy as np
import math
from scipy.spatial import cKDTree
import gdal
import osr
import subprocess
from pdb import set_trace
import tempfile
import argparse
import os
from glob import glob
import tempfile
import parasol
from parasol import lidar, common, cfg
logger = logging.getLogger(__name__)
# local constants
PAD = 10 # meters
TILE_DIM = 1000 # meters
def grid_points(x_min, x_max, y_min, y_max):
"""
Grid scattered points using kNN median filter
Arguments:
x_min, x_max, y_min, y_max: floats, limits for bounding box
grnd: set True to compute surface from ground returns, False to compute
surface from (first) non-ground returns
Returns: x_vec, y_vec, z_grd
x_vec, y_vec: numpy 1D arrays, x and y coordinate axes
z_surf, z_grnd: numpy 2D arrays, elevation grids
"""
# Note from LiDAR metadata: ... Default (Class 1), Ground (Class 2), Noise
# (Class 7), Water (Class 9), Ignored Ground (Class 10), Overlap Default
# (Class 17) and Overlap Ground (Class 18).
# build output grid spanning bbox
x_vec = np.arange(math.floor(x_min), math.floor(x_max), cfg.SURFACE_RES_M)
y_vec = np.arange(math.floor(y_min), math.floor(y_max), cfg.SURFACE_RES_M)
x_grd, y_grd = np.meshgrid(x_vec, y_vec)
# retrieve data, including a pad on all sides
pts = lidar.retrieve(x_min-PAD, x_max+PAD, y_min-PAD, y_max+PAD)
# extract ground points
grnd_idx = []
for idx, pt in enumerate(pts):
if pt[3] == pt[4] and pt[5] in {1, 2, 9}:
# last or only return, classified as "default", "ground" or "water"
grnd_idx.append(idx)
grnd_pts = pts[grnd_idx, :3]
# extract upper surface points
surf_idx = []
for idx, pt in enumerate(pts):
if (pt[3] == 1 or pt[4] == 1) and pt[5] in {1, 2, 9}:
# first or only return, classified as "default", "ground", or "water"
surf_idx.append(idx)
surf_pts = pts[surf_idx, :3]
del pts
z_grds = []
for pts in [grnd_pts, surf_pts]:
# extract [x, y] and z arrays
xy = pts[:, :2]
zz = pts[:, 2]
# find NN for all grid points
tree = cKDTree(xy)
xy_grd = np.hstack([x_grd.reshape((-1,1)), y_grd.reshape((-1,1))])
nn_dist, nn_idx = tree.query(xy_grd, k=16)
# compute local medians
z_grds.append(np.median(zz[nn_idx], axis=1).reshape(x_grd.shape))
return x_vec, y_vec, z_grds[0], z_grds[1]
def create_geotiff(filename, x_vec, y_vec, z_grd):
"""
Write input array as GeoTiff raster
Arguments:
filename: string, path to write GeoTiff file
x_vec, y_vec: numpy 1D arrays, x and y coordinate axes
z_grd: numpy 2D array, elevation grid
Returns: Nothing, writes result to file
"""
# invert y-axis
y_vec = y_vec[::-1]
z_grd = z_grd[::-1,:]
# create file
rows, cols = z_grd.shape
driver = gdal.GetDriverByName('GTiff')
outRaster = driver.Create(filename, cols, rows, 1, gdal.GDT_Float32)
outRaster.SetGeoTransform((x_vec[0], cfg.SURFACE_RES_M, 0, y_vec[0], 0, -cfg.SURFACE_RES_M))
outband = outRaster.GetRasterBand(1)
outband.WriteArray(z_grd)
outRasterSRS = osr.SpatialReference()
outRasterSRS.ImportFromEPSG(cfg.PRJ_SRID)
outRaster.SetProjection(outRasterSRS.ExportToWkt())
outband.FlushCache()
# clean up
driver = outRaster = outband = None
# TODO: make this concurrent
def create_surfaces(x_min, x_max, y_min, y_max, x_tile, y_tile):
"""
Generate rasters and upload to database, tile-by-tile
Arguments:
x_min, x_max, y_min, y_max: floats, limits for the full region-of-interest
x_tile, y_tile: floats, desired dimensions for generated tiles, note that
the actual dimensions are adjusted to evenly divide the ROI
Returns: nothing
"""
tiles = common.tile_limits(x_min, x_max, y_min, y_max, x_tile, y_tile)
num_tiles = len(tiles)
modes = ['add'] * len(tiles)
modes[0] = 'create'
modes[-1] = 'finish'
# create directory, if needed
if not os.path.isdir(cfg.SURFACE_DIR):
os.makedirs(cfg.SURFACE_DIR)
# generate tiles
for ii, tile in enumerate(tiles):
logger.info(f'Generating tile {ii+1} of {num_tiles}')
x_vec, y_vec, z_grnd, z_surf = grid_points(**tile)
logger.info(f'Gridding ground points')
grnd_name = os.path.join(cfg.SURFACE_DIR, 'ground_tile_{:04d}.tif'.format(ii))
create_geotiff(grnd_name, x_vec, y_vec, z_grnd)
logger.info(f'Wrote tile {grnd_name}')
logger.info(f'Gridding surface points')
surf_name = os.path.join(cfg.SURFACE_DIR, 'surface_tile_{:04d}.tif'.format(ii))
create_geotiff(surf_name, x_vec, y_vec, z_surf)
logger.info(f'Wrote tile {surf_name}')
# merge tiles
logger.info('Merging tiles')
for prefix in ['ground', 'surface']:
cmd = ['gdal_merge.py']
cmd.extend(glob(os.path.join(cfg.SURFACE_DIR, prefix + '_tile_*.tif')))
cmd.extend(['-of', 'GTiff', '-o', os.path.join(cfg.SURFACE_DIR, prefix + '.tif')])
subprocess.run(cmd)
# delete tiles, which are no longer needed
logger.info('Deleting tile files')
for fn in glob(os.path.join(cfg.SURFACE_DIR, '*_tile_*.tif')):
os.remove(fn)
def retrieve(x_min, x_max, y_min, y_max, which):
"""
Retrieve subset within specified ROI
Arguments:
minx, maxx, miny, maxy: floats, limits for bounding box
which: string, which raster to retrieve data from, must be one of
'surface', 'ground'
Returns: x_grd, y_grd, z_grd
"""
# get raster file name
if which == 'surface':
src_file = os.path.join(cfg.SURFACE_DIR, 'surface.tif')
elif which == 'ground':
src_file = os.path.join(cfg.SURFACE_DIR, 'ground.tif')
else:
raise ValueError('Invalid choice for "which" variable')
# get subset as file
with tempfile.NamedTemporaryFile() as fp:
dest_file = fp.name
subprocess.run(['gdal_translate', '-of', 'GTiff', '-projwin',
str(x_min), str(y_max), str(x_max), str(y_min), src_file, dest_file])
ds = gdal.Open(dest_file)
z_grd = np.array(ds.GetRasterBand(1).ReadAsArray())
transform = ds.GetGeoTransform()
y_pixel = np.arange(0, z_grd.shape[0]).reshape((z_grd.shape[0], 1))
x_pixel = np.arange(0, z_grd.shape[1]).reshape((1, z_grd.shape[1]))
x_grd = transform[0] + x_pixel*transform[1] + y_pixel*transform[2]
y_grd = transform[3] + x_pixel*transform[4] + y_pixel*transform[5]
return x_grd, y_grd, z_grd
# command line utilities -----------------------------------------------------
def initialize_cli():
"""Command line utility for initializing the surface/ground databases"""
ap = argparse.ArgumentParser(
description="Initialize Parasol surface & ground raster databases - clobbers existing!",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
ap.add_argument('--log', type=str, default='info', help="select logging level",
choices=['debug', 'info', 'warning', 'error', 'critical'])
ap.add_argument('--dryrun', action='store_true', help='Set to preview only')
args = ap.parse_args()
log_lvl = getattr(logging, args.log.upper())
logging.basicConfig(level=log_lvl)
logger.setLevel(log_lvl)
if (args.dryrun):
tiles = common.tile_limits(cfg.DOMAIN_XLIM[0], cfg.DOMAIN_XLIM[1],
cfg.DOMAIN_YLIM[0], cfg.DOMAIN_YLIM[1], TILE_DIM, TILE_DIM)
print(f'Compute raster in domain: [[{cfg.DOMAIN_XLIM}], [{cfg.DOMAIN_YLIM}]]')
print(f'Num tiles: {len(tiles)}')
else:
create_surfaces(cfg.DOMAIN_XLIM[0], cfg.DOMAIN_XLIM[1], cfg.DOMAIN_YLIM[0],
cfg. DOMAIN_YLIM[1], TILE_DIM, TILE_DIM)