From a4b586274027308ea68fb74ca0e0bc20522936e4 Mon Sep 17 00:00:00 2001 From: Piero Toffanin Date: Mon, 28 Oct 2019 14:40:40 -0400 Subject: [PATCH] Faster orthophoto merging --- opendm/orthophoto.py | 95 ++++++++++++++++++++++++++++++---------- opendm/remote.py | 1 + stages/odm_orthophoto.py | 8 +++- stages/splitmerge.py | 60 +++---------------------- 4 files changed, 85 insertions(+), 79 deletions(-) diff --git a/opendm/orthophoto.py b/opendm/orthophoto.py index 988ce24b..84c2ada5 100644 --- a/opendm/orthophoto.py +++ b/opendm/orthophoto.py @@ -6,7 +6,10 @@ from opendm.concurrency import get_max_memory import math import numpy as np import rasterio +import fiona +from scipy import ndimage from rasterio.transform import Affine, rowcol +from rasterio.mask import mask from opendm import io def get_orthophoto_vars(args): @@ -49,7 +52,49 @@ def post_orthophoto_steps(args, bounds_file_path, orthophoto_file): if args.orthophoto_png: generate_png(orthophoto_file) -def merge(input_ortho_and_cutlines, output_orthophoto, blend_distance=20, orthophoto_vars={}): + +def compute_mask_raster(input_raster, vector_mask, output_raster, blend_distance=20, only_max_coords_feature=False): + if not os.path.exists(input_raster): + log.ODM_WARNING("Cannot mask raster, %s does not exist" % input_raster) + return + + if not os.path.exists(vector_mask): + log.ODM_WARNING("Cannot mask raster, %s does not exist" % vector_mask) + return + + with rasterio.open(input_raster, 'r') as rast: + with fiona.open(vector_mask) as src: + burn_features = src + + if only_max_coords_feature: + max_coords_count = 0 + max_coords_feature = None + for feature in src: + if feature is not None: + # No complex shapes + if len(feature['geometry']['coordinates'][0]) > max_coords_count: + max_coords_count = len(feature['geometry']['coordinates'][0]) + max_coords_feature = feature + if max_coords_feature is not None: + burn_features = [max_coords_feature] + + shapes = [feature["geometry"] for feature in burn_features] + out_image, out_transform = mask(rast, shapes, nodata=0) + + if blend_distance > 0: + alpha_band = out_image[3] + dist_t = ndimage.distance_transform_edt(alpha_band) + dist_t[dist_t <= blend_distance] /= blend_distance + dist_t[dist_t > blend_distance] = 1 + np.multiply(alpha_band, dist_t, out=alpha_band, casting="unsafe") + + with rasterio.open(output_raster, 'w', **rast.profile) as dst: + dst.write(out_image) + + return output_raster + + +def merge(input_ortho_and_ortho_cuts, output_orthophoto, orthophoto_vars={}): """ Based on https://github.com/mapbox/rio-merge-rgba/ Merge orthophotos around cutlines using a blend buffer. @@ -58,7 +103,7 @@ def merge(input_ortho_and_cutlines, output_orthophoto, blend_distance=20, orthop bounds=None precision=7 - for o, c in input_ortho_and_cutlines: + for o, c in input_ortho_and_ortho_cuts: if not io.file_exists(o): log.ODM_WARNING("%s does not exist. Will skip from merged orthophoto." % o) continue @@ -72,24 +117,23 @@ def merge(input_ortho_and_cutlines, output_orthophoto, blend_distance=20, orthop return with rasterio.open(inputs[0][0]) as first: - src_nodata = first.nodatavals[0] # TODO: this is None res = first.res dtype = first.dtypes[0] profile = first.profile log.ODM_INFO("%s valid orthophoto rasters to merge" % len(inputs)) - sources = [(rasterio.open(o)) for o,_ in inputs] + sources = [(rasterio.open(o), rasterio.open(c)) for o,c in inputs] # scan input files. # while we're at it, validate assumptions about inputs xs = [] ys = [] - for src in sources: + for src, _ in sources: left, bottom, right, top = src.bounds xs.extend([left, right]) ys.extend([bottom, top]) - # if src.profile["count"] != 1 or src.profile["count"] != 1: - # raise ValueError("Inputs must be 1-band rasters") + if src.profile["count"] != 4: + raise ValueError("Inputs must be 4-band rasters") dst_w, dst_s, dst_e, dst_n = min(xs), min(ys), max(xs), max(ys) log.ODM_INFO("Output bounds: %r %r %r %r" % (dst_w, dst_s, dst_e, dst_n)) @@ -115,7 +159,6 @@ def merge(input_ortho_and_cutlines, output_orthophoto, blend_distance=20, orthop profile["compress"] = orthophoto_vars.get('COMPRESS', 'LZW') profile["predictor"] = orthophoto_vars.get('PREDICTOR', '2') profile["bigtiff"] = orthophoto_vars.get('BIGTIFF', 'IF_SAFER') - profile["nodata"] = src_nodata profile.update() # create destination file @@ -130,23 +173,10 @@ def merge(input_ortho_and_cutlines, output_orthophoto, blend_distance=20, orthop dst_count = first.count dst_shape = (dst_count, dst_rows, dst_cols) + # First pass, write all rasters naively dstarr = np.zeros(dst_shape, dtype=dtype) - distsum = np.zeros(dst_shape, dtype=dtype) - for src in sources: - # The full_cover behavior is problematic here as it includes - # extra pixels along the bottom right when the sources are - # slightly misaligned - # - # src_window = get_window(left, bottom, right, top, - # src.transform, precision=precision) - # - # With rio merge this just adds an extra row, but when the - # imprecision occurs at each block, you get artifacts - - nodata = src.nodatavals[0] - - # Alternative, custom get_window using rounding + for src, _ in sources: src_window = tuple(zip(rowcol( src.transform, left, top, op=round, precision=precision ), rowcol( @@ -168,6 +198,25 @@ def merge(input_ortho_and_cutlines, output_orthophoto, blend_distance=20, orthop if np.count_nonzero(dstarr[3]) == blocksize: break + # Second pass, write cut rasters + for _, cut in sources: + src_window = tuple(zip(rowcol( + cut.transform, left, top, op=round, precision=precision + ), rowcol( + cut.transform, right, bottom, op=round, precision=precision + ))) + + temp = np.zeros(dst_shape, dtype=dtype) + temp = cut.read( + out=temp, window=src_window, boundless=True, masked=False + ) + + # For each band, average alpha values between + # destination raster and cut raster + for b in range(0, 3): + blended = temp[3] / 255.0 * temp[b] + (1 - temp[3] / 255.0) * dstarr[b] + np.copyto(dstarr[b], blended, casting='unsafe', where=temp[3]!=0) + dstrast.write(dstarr, window=dst_window) return output_orthophoto diff --git a/opendm/remote.py b/opendm/remote.py index 04657201..74000d17 100644 --- a/opendm/remote.py +++ b/opendm/remote.py @@ -501,6 +501,7 @@ class ToolchainTask(Task): "opensfm/exif/empty"], outputs=["odm_orthophoto/odm_orthophoto.tif", "odm_orthophoto/cutline.gpkg", + "odm_orthophoto/odm_orthophoto_cut.tif", "odm_dem", "odm_georeferencing"]) else: diff --git a/stages/odm_orthophoto.py b/stages/odm_orthophoto.py index 2dfc59ce..9321632d 100644 --- a/stages/odm_orthophoto.py +++ b/stages/odm_orthophoto.py @@ -108,13 +108,19 @@ class ODMOrthoPhotoStage(types.ODM_Stage): # Cutline computation, before cropping # We want to use the full orthophoto, not the cropped one. if args.orthophoto_cutline: + cutline_file = os.path.join(tree.odm_orthophoto, "cutline.gpkg") + compute_cutline(tree.odm_orthophoto_tif, bounds_file_path, - os.path.join(tree.odm_orthophoto, "cutline.gpkg"), + cutline_file, args.max_concurrency, tmpdir=os.path.join(tree.odm_orthophoto, "grass_cutline_tmpdir"), scale=0.25) + orthophoto.compute_mask_raster(tree.odm_orthophoto_tif, cutline_file, + os.path.join(tree.odm_orthophoto, "odm_orthophoto_cut.tif"), + blend_distance=20, only_max_coords_feature=True) + orthophoto.post_orthophoto_steps(args, bounds_file_path, tree.odm_orthophoto_tif) geotiffcreated = True diff --git a/stages/splitmerge.py b/stages/splitmerge.py index bbacd552..b4d4bad5 100644 --- a/stages/splitmerge.py +++ b/stages/splitmerge.py @@ -217,72 +217,22 @@ class ODMMergeStage(types.ODM_Stage): system.mkdir_p(tree.odm_orthophoto) if not io.file_exists(tree.odm_orthophoto_tif) or self.rerun(): - all_orthos_and_cutlines = get_all_submodel_paths(tree.submodels_path, + all_orthos_and_ortho_cuts = get_all_submodel_paths(tree.submodels_path, os.path.join("odm_orthophoto", "odm_orthophoto.tif"), - os.path.join("odm_orthophoto", "cutline.gpkg"), + os.path.join("odm_orthophoto", "odm_orthophoto_cut.tif"), ) - if len(all_orthos_and_cutlines) > 1: + if len(all_orthos_and_ortho_cuts) > 1: log.ODM_INFO("Found %s submodels with valid orthophotos and cutlines" % len(all_orthos_and_cutlines)) # TODO: histogram matching via rasterio # currently parts have different color tones - merged_geotiff = os.path.join(tree.odm_orthophoto, "odm_orthophoto.merged.tif") - - kwargs = { - 'orthophoto_merged': merged_geotiff, - 'input_files': ' '.join(map(lambda i: quote(i[0]), all_orthos_and_cutlines)), - 'max_memory': get_max_memory(), - 'threads': args.max_concurrency, - } - - # use bounds as cutlines (blending) - if io.file_exists(merged_geotiff): - os.remove(merged_geotiff) - - system.run('gdal_merge.py -o {orthophoto_merged} ' - #'-createonly ' - '-co "BIGTIFF=YES" ' - '-co "BLOCKXSIZE=512" ' - '-co "BLOCKYSIZE=512" ' - '--config GDAL_CACHEMAX {max_memory}% ' - '{input_files} '.format(**kwargs) - ) - - for ortho_cutline in all_orthos_and_cutlines: - kwargs['input_file'], kwargs['cutline'] = ortho_cutline - - # Note: cblend has a high performance penalty - system.run('gdalwarp -cutline {cutline} ' - '-cblend 20 ' - '-r bilinear -multi ' - '-wo NUM_THREADS={threads} ' - '--config GDAL_CACHEMAX {max_memory}% ' - '{input_file} {orthophoto_merged}'.format(**kwargs) - ) - - # Apply orthophoto settings (compression, tiling, etc.) - orthophoto_vars = orthophoto.get_orthophoto_vars(args) - if io.file_exists(tree.odm_orthophoto_tif): os.remove(tree.odm_orthophoto_tif) - kwargs = { - 'vars': ' '.join(['-co %s=%s' % (k, orthophoto_vars[k]) for k in orthophoto_vars]), - 'max_memory': get_max_memory(), - 'merged': merged_geotiff, - 'log': tree.odm_orthophoto_tif_log, - 'orthophoto': tree.odm_orthophoto_tif, - } - - system.run('gdal_translate ' - '{vars} ' - '--config GDAL_CACHEMAX {max_memory}% ' - '{merged} {orthophoto} > {log}'.format(**kwargs)) - - os.remove(merged_geotiff) - + orthophoto_vars = orthophoto.get_orthophoto_vars(args) + orthophoto.merge(all_orthos_and_ortho_cuts, tree.odm_orthophoto_tif, orthophoto_vars) orthophoto.post_orthophoto_steps(args, merged_bounds_file, tree.odm_orthophoto_tif) elif len(all_orthos_and_cutlines) == 1: # Simply copy