def get_branch_mask(x_y, mgrid_n_xy):
"""Obtains total coverage mask of single branch"""
# Obtain mask(s) of the root point(s)
res_mask = numpy.zeros(mgrid_n_xy.shape[:-1], dtype=bool)
gdal_utils.write_arr(res_mask, x_y, True)
mask = gdal_utils.read_arr(res_mask, mgrid_n_xy)
# The AND-NOT is needed to drop the self-pointing graph-seeds
mask &= ~res_mask
while mask.any():
res_mask |= mask
mask = gdal_utils.read_arr(mask, mgrid_n_xy)
return res_mask
def process_neighbors(dem_band, dir_arr, x_y):
"""Process the valid and pending neighbor points and return a list to be put to tentative"""
x_y = x_y[..., numpy.newaxis, :]
n_xy = neighbor_xy(x_y, VALID_NEIGHBOR_DIRS)
n_dir = numpy.broadcast_to(VALID_NEIGHBOR_DIRS, n_xy.shape[:-1])
# Filter out of bounds pixels
mask = dem_band.in_bounds(n_xy)
if not mask.all():
n_xy = n_xy[mask]
n_dir = n_dir[mask]
# The lines can only pass-thru inner DEM pixels, the boundary ones do split
stop_mask = ~mask.all(-1)
# Filter already processed pixels
neighs = gdal_utils.read_arr(dir_arr, n_xy)
mask = neighs == NEIGHBOR_PENDING
if not mask.any():
return None
if not mask.all():
n_xy = n_xy[mask]
n_dir = n_dir[mask]
mask = neighs == NEIGHBOR_BOUNDARY
if mask.any():
stop_mask |= mask.any(-1)
# Process selected pixels
n_dir = neighbor_flip(n_dir)
# Put 'stop' markers on the successors of the masked points
# This is to split lines at the boundary pixels
if stop_mask.any():
n_dir[stop_mask] = NEIGHBOR_STOP
gdal_utils.write_arr(dir_arr, n_xy, n_dir)
return n_xy
def flip_lines(dir_arr, x_y):
"""Flip all 'n_dir'-s along multiple lines"""
prev_dir = gdal_utils.read_arr(dir_arr, x_y)
gdal_utils.write_arr(dir_arr, x_y, NEIGHBOR_STOP)
while True:
n_xy = neighbor_xy(x_y, prev_dir)
n_dir = gdal_utils.read_arr(dir_arr, n_xy)
gdal_utils.write_arr(dir_arr, n_xy, neighbor_flip(prev_dir))
mask = neighbor_is_invalid(n_dir)
if mask.any():
assert ((n_dir[mask] == NEIGHBOR_SEED) |
(n_dir[mask] == NEIGHBOR_STOP)).all()
if mask.all():
return dir_arr
mask = ~mask
n_xy = n_xy[mask]
n_dir = n_dir[mask]
x_y = n_xy
prev_dir = n_dir
def main(argv):
"""Main entry"""
valleys = False
boundary_val = None
truncate = True
src_filename = dst_filename = None
while argv:
if argv[0][0] == '-':
if argv[0] == '-h':
return print_help()
if argv[0] == '-valley':
valleys = True
elif argv[0] == '-boundary_val':
argv = argv[1:]
boundary_val = float(argv[0])
else:
return print_help('Unsupported option "%s"' % argv[0])
else:
if src_filename is None:
src_filename = argv[0]
elif dst_filename is None:
dst_filename = argv[0]
else:
return print_help('Unexpected argument "%s"' % argv[0])
argv = argv[1:]
if src_filename is None or dst_filename is None:
return print_help('Missing file-names')
# Load DEM
dem_band = gdal_utils.dem_open(src_filename)
if dem_band is None:
return print_help('Unable to open "%s"' % src_filename)
dst_ds = gdal_utils.vect_create(dst_filename)
if dst_ds is None:
return print_help('Unable to create "%s"' % src_filename)
dem_band.load()
#
# Trace ridges/valleys
#
if RESUME_FROM_SNAPSHOT < 1:
start = time.time()
# Actual trace
dir_arr = trace_ridges(dem_band, valleys, boundary_val)
if dir_arr is None:
print('Error: Failed to trace ridges', file=sys.stderr)
return 2
duration = time.time() - start
print('Traced through %d/%d points, %d sec' %
(numpy.count_nonzero(~neighbor_is_invalid(dir_arr)),
dir_arr.size, duration))
if KEEP_SNAPSHOT:
keep_arrays(src_filename + '-1-', {
'dir_arr': dir_arr,
})
elif RESUME_FROM_SNAPSHOT == 1:
dir_arr, = restore_arrays(src_filename + '-1-', {
'dir_arr': None,
})
#
# The coverage-area of each pixels is needed by arrange_lines()
# The distance object is used to calculate the branch length
#
distance = gdal_utils.geod_distance(dem_band) if 0 == DISTANCE_METHOD \
else gdal_utils.tm_distance(dem_band) if 1 == DISTANCE_METHOD \
else gdal_utils.draft_distance(dem_band)
area_arr = calc_pixel_area(distance, dem_band.shape)
print('Calculated total area %.2f km2, mean %.2f m2' %
(area_arr.sum() / 1e6, area_arr.mean()))
#
# Identify and flip the "trunk" branches
# All the real-seeds become regular graph-nodes or "leaf" pixel.
# The former start/leaf pixel of these branches becomes a "seed".
#
if RESUME_FROM_SNAPSHOT < 2:
start = time.time()
# Arrange branches to select which one to flip (trunks_only)
branch_lines = arrange_lines(dir_arr, area_arr, True)
# Actual flip
if flip_lines(dir_arr, branch_lines['start_xy']) is None:
print('Error: Failed to flip %d branches' % (branch_lines.size),
file=sys.stderr)
return 2
duration = time.time() - start
print(
'Flip & merge total %d trunk-branches, max/min area %.1f/%.3f km2, %d sec'
% (branch_lines.size, branch_lines['area'].max() / 1e6,
branch_lines['area'].min() / 1e6, duration))
if KEEP_SNAPSHOT:
keep_arrays(src_filename + '-2-', {
'dir_arr': dir_arr,
'branch_lines': branch_lines,
})
elif RESUME_FROM_SNAPSHOT == 2:
dir_arr, branch_lines = restore_arrays(
src_filename + '-2-', {
'dir_arr': None,
'branch_lines': BRANCH_LINE_DTYPE,
})
#
# Identify all the branches
#
if RESUME_FROM_SNAPSHOT < 3:
start = time.time()
# Arrange branches
branch_lines = arrange_lines(dir_arr, area_arr, False)
# Sort the the generated branches (descending 'area' order)
argsort = numpy.argsort(branch_lines['area'])
branch_lines = numpy.take(branch_lines, argsort[::-1])
# Trim to 5 zoom-levels (1/1024 of max area)
min_area = area_arr.sum() / (4**5)
print(
' Trimming total %d branches to min area of %.3f km2 (currently %.3f km2)'
% (branch_lines.size, min_area / 1e6,
branch_lines['area'].min() / 1e6))
branch_lines = branch_lines[branch_lines['area'] >= min_area]
duration = time.time() - start
print('Created total %d branches, max/min area %.1f/%.3f km2, %d sec' %
(branch_lines.size, branch_lines['area'].max() / 1e6,
branch_lines['area'].min() / 1e6, duration))
if KEEP_SNAPSHOT:
keep_arrays(src_filename + '-3-', {
'branch_lines': branch_lines,
})
elif RESUME_FROM_SNAPSHOT == 3:
dir_arr, = restore_arrays(src_filename + '-2-', {
'dir_arr': None,
})
branch_lines, = restore_arrays(src_filename + '-3-', {
'branch_lines': BRANCH_LINE_DTYPE,
})
if dst_ds:
start = time.time()
# Delete existing layers
if truncate:
for i in reversed(range(dst_ds.get_layer_count())):
print(' Deleting layer',
gdal_utils.gdal_vect_layer(dst_ds, i).get_name())
dst_ds.delete_layer(i)
# Create new one
layer_options = DEF_LAYER_OPTIONS
bydrv_options = BYDVR_LAYER_OPTIONS.get(dst_ds.get_drv_name())
if bydrv_options:
layer_options += bydrv_options
dst_layer = gdal_utils.gdal_vect_layer.create(
dst_ds,
VECTOR_LAYER_NAME(valleys),
srs=dem_band.get_spatial_ref(),
geom_type=gdal_utils.wkbLineString,
options=layer_options)
if dst_layer is None:
print('Error: Unable to create layer', file=sys.stderr)
return 1
# Add fields
dst_layer.create_field('Name', gdal_utils.OFTString) # KML <name>
dst_layer.create_field('Description',
gdal_utils.OFTString) # KML <description>
if FEATURE_OSM_NATURAL:
dst_layer.create_field('natural',
gdal_utils.OFTString) # OSM "natural" key
# Note that mgrid_n_xy is for coverage area assert only
mgrid_n_xy = neighbor_xy_safe(get_mgrid(dir_arr.shape), dir_arr)
geometries = 0
for branch in branch_lines:
ar = calc_branch_area(branch['x_y'], mgrid_n_xy, area_arr)
assert int(branch['area']) == int(
ar
), 'Accumulated branch coverage area mismatch %.6f / %.6f km2' % (
branch['area'] / 1e6, ar / 1e6)
# Advance one step forward to connect to the parent branch
if not SEPARATED_BRANCHES:
x_y = branch['x_y']
branch['x_y'] = neighbor_xy_safe(
x_y, gdal_utils.read_arr(dir_arr, x_y))
# Extract the branch pixel coordinates and calculate length
x_y = branch['start_xy']
polyline = [x_y]
dist = 0.
while (x_y != branch['x_y']).any():
# Advance to the next point
new_xy = neighbor_xy(x_y, gdal_utils.read_arr(dir_arr, x_y))
dist += distance.get_distance(x_y, new_xy)
x_y = new_xy
polyline.append(x_y)
# Create actual geometry
geom = dst_layer.create_feature_geometry(gdal_utils.wkbLineString)
geom.set_field(
'Name',
'%dm' % dist if dist < 10000 else '%dkm' % round(dist / 1000))
geom.set_field(
'Description', 'length: %.1f km, area: %.1f km2' %
(dist / 1e3, branch['area'] / 1e6))
if FEATURE_OSM_NATURAL:
geom.set_field('natural', FEATURE_OSM_NATURAL(valleys))
geom.set_style_string(VECTOR_FEATURE_STYLE(valleys))
# Reverse the line to match the tracing direction
for x_y in reversed(polyline):
geom.add_point(*dem_band.xy2lonlatalt(x_y))
geom.create()
geometries += 1
duration = time.time() - start
print('Created total %d geometries, %d sec' % (geometries, duration))
return 0
def arrange_lines(dir_arr, area_arr, trunks_only):
"""Arrange lines in branches by using the area of coverage"""
area_arr = area_arr.copy()
# Count the number of neighbors pointing to each pixel
# Only the last branch that reaches a pixel continues forward, others stop there.
# As the branches are processed starting from the one with less coverage-area, this
# allows the largest one to reach the graph-seed.
mgrid_n_xy = neighbor_xy_safe(get_mgrid(dir_arr.shape), dir_arr)
n_num = numpy.zeros(dir_arr.shape, dtype=int)
for d in VALID_NEIGHBOR_DIRS:
n_xy = mgrid_n_xy[dir_arr == d]
n = numpy.zeros_like(n_num)
gdal_utils.write_arr(n, n_xy, 1)
n_num += n
del n_xy, n
# Put -1 at invalid nodes, except the "real" seeds (distinguish from the "leafs")
valid_mask = ~neighbor_is_invalid(dir_arr)
n_num[~valid_mask & (n_num == 0)] = -1
all_leafs = n_num == 0
print('Detected %d "leaf" and %d "real-seed" pixels' %
(numpy.count_nonzero(all_leafs),
numpy.count_nonzero(~valid_mask & (n_num > 0))))
# Start at the "leaf" pixels
pend_lines = numpy.zeros(numpy.count_nonzero(all_leafs),
dtype=BRANCH_LINE_DTYPE)
pend_lines['start_xy'] = pend_lines['x_y'] = numpy.array(
numpy.nonzero(all_leafs)).T
pend_lines['area'] = gdal_utils.read_arr(area_arr, pend_lines['x_y'])
#
# Process the leaf-branches in parallel
# The parallel processing must stop at the point before the graph-nodes
#
bridge_mask = gdal_utils.read_arr(n_num == 1, mgrid_n_xy)
bridge_mask &= valid_mask
all_leafs[...] = False
pend_mask = numpy.ones(pend_lines.size, dtype=bool)
x_y = pend_lines['x_y']
while pend_mask.any():
gdal_utils.write_arr(all_leafs, x_y, True)
# Stop in front the graph nodes and at the "seeds"
mask = gdal_utils.read_arr(bridge_mask, x_y)
x_y = x_y[mask]
# Advance the points, which are still at graph-bridges
x_y = gdal_utils.read_arr(mgrid_n_xy, x_y)
assert (gdal_utils.read_arr(n_num, x_y) == 1).all()
gdal_utils.write_arr(n_num, x_y, 0)
# Keep the intermediate results
pend_mask[pend_mask] = mask
pend_lines['x_y'][pend_mask] = x_y
# Accumulate coverage area
area = gdal_utils.read_arr(area_arr, x_y)
pend_lines['area'][pend_mask] += area
del bridge_mask
del pend_mask
assert int(pend_lines['area'].sum()) == int(area_arr[all_leafs].sum(
)), 'Leaf-branch coverage area mismatch %.6f / %.6f km2' % (
pend_lines['area'].sum() / 1e6, area_arr[all_leafs].sum() / 1e6)
# Trim leaf-trunks
mask = gdal_utils.read_arr(valid_mask, pend_lines['x_y'])
pend_lines = pend_lines[mask]
print(
' Detected %d pixels in "leaf" branches, area %.2f km2, trim %d leaf-trunks'
% (numpy.count_nonzero(all_leafs), pend_lines['area'].sum() / 1e6,
numpy.count_nonzero(~mask)))
# Update the accumulated area, but only at the stop-points (in front the graph-nodes)
gdal_utils.write_arr(area_arr, pend_lines['x_y'], pend_lines['area'])
branch_lines = numpy.empty_like(pend_lines, shape=[0])
trim_cnt = 0
progress_idx = 0
while pend_lines.size:
# Process the branch with minimal coverage-area
br_idx = pend_lines['area'].argmin()
branch = pend_lines[br_idx]
x_y = branch['x_y']
area = gdal_utils.read_arr(area_arr, x_y)
assert branch[
'area'] <= area, 'Branch area decreases at %s: %f -> %f m2' % (
x_y, branch['area'], area)
if gdal_utils.read_arr(valid_mask, x_y):
# Advance to the next point
x_y = gdal_utils.read_arr(mgrid_n_xy, x_y)
# Accumulate the coverage-area
area += gdal_utils.read_arr(area_arr, x_y)
gdal_utils.write_arr(area_arr, x_y, area)
# Handle node-bridges counter: only the last branch to proceed further
n = gdal_utils.read_arr(n_num, x_y)
assert n > 0
gdal_utils.write_arr(n_num, x_y, n - 1)
# Stop at graph-node (non-last branches)
is_stop = n > 1
keep_branch = not trunks_only
# Update the end-point
if not is_stop:
branch['x_y'] = x_y
branch['area'] = area
else:
# Stop at graph-seed (trunk branch)
keep_branch = is_stop = True
if is_stop:
# Discard the "leaf" branches, with "trunks_only" -- non-trunk branches
if keep_branch:
if False == gdal_utils.read_arr(all_leafs, branch['x_y']):
branch_lines = numpy.append(branch_lines, branch)
else:
trim_cnt += 1
pend_lines = numpy.delete(pend_lines, br_idx)
#
# Progress, each 10000-th step
#
if progress_idx % 10000 == 0:
area = branch_lines['area'].max() if branch_lines.size else 0
if pend_lines.size:
area = max(area, pend_lines['area'].max())
print(
' Process step %d, max. area %.2f km2, completed %d, pending %d, trimmed leaves %d'
% (progress_idx, area / 1e6, branch_lines.size,
pend_lines.size, trim_cnt))
progress_idx += 1
# Confirm everything is processed
assert (n_num <= 0).all(), 'Unprocessed pixels at %s' % numpy.array(
numpy.nonzero(n_num > 0)).T
return branch_lines