diff --git a/opendm/orthophoto.py b/opendm/orthophoto.py index 27f42b81a..628b6b134 100644 --- a/opendm/orthophoto.py +++ b/opendm/orthophoto.py @@ -1,4 +1,8 @@ import os +import threading +import contextlib +from collections import deque +from concurrent.futures import ThreadPoolExecutor from opendm import log from opendm import system from opendm.cropper import Cropper @@ -18,6 +22,22 @@ from osgeo import ogr +@contextlib.contextmanager +def _bounded_gdal_cache(nbytes): + """Temporarily cap GDAL's global block cache, restoring it on exit. + + A small cache keeps the parallel merge's output-tile flushes prompt and + cheap. Restoring on the way out (even if the merge raises) avoids leaving + the rest of the pipeline (notably the COG conversion) with a shrunken cache. + """ + prev = gdal.GetCacheMax() + gdal.SetCacheMax(nbytes) + try: + yield + finally: + gdal.SetCacheMax(prev) + + def get_orthophoto_vars(args): return { 'TILED': 'NO' if args.orthophoto_no_tiled else 'YES', @@ -265,10 +285,55 @@ def feather_raster(input_raster, output_raster, blend_distance=20): return output_raster -def merge(input_ortho_and_ortho_cuts, output_orthophoto, orthophoto_vars={}, merge_skip_blending=False): +def _read_window_gated(ds, src_window, dst_shape, dtype): + """Read ``src_window`` into a ``dst_shape`` array, equivalent to a + ``boundless=True`` read with 0 nodata fill, but avoiding rasterio's boundless + VRT path for the common cases. ``boundless=True`` builds a VRT and serializes + it via Python's ElementTree (``_serialize_xml``) on every read, which is a + large per-read overhead on big merges. Behaviour: + - window fully outside the dataset -> all zeros (no read), matching the 0 + nodata fill a boundless read would produce here; + - window fully inside -> plain non-boundless read (no VRT), identical to a + boundless read when no out-of-bounds padding is needed; + - window partially overlapping the edge -> fall back to boundless (rare; + only the true border blocks), preserving exact fill behaviour. + """ + (r0, r1), (c0, c1) = src_window + out = np.zeros(dst_shape, dtype=dtype) + height, width = ds.height, ds.width + if r1 <= 0 or c1 <= 0 or r0 >= height or c0 >= width: + return out + if r0 >= 0 and c0 >= 0 and r1 <= height and c1 <= width: + return ds.read(out=out, window=src_window, boundless=False, masked=False) + return ds.read(out=out, window=src_window, boundless=True, masked=False) + + +# Hard cap on orthophoto-merge worker threads. Each worker opens every input pair +# (ortho + cut), so open file handles scale as 2 * pairs * max_workers. With many +# submodels on a high-core machine that can exhaust the open-file limit (EMFILE). +# A fixed cap keeps a meaningful speedup while staying well under typical limits, and +# is simpler and more portable than probing RLIMIT_NOFILE (absent on Windows). +MERGE_WORKER_CAP = 16 + +def merge(input_ortho_and_ortho_cuts, output_orthophoto, orthophoto_vars={}, merge_skip_blending=False, max_workers=1): """ Based on https://github.com/mapbox/rio-merge-rgba/ Merge orthophotos around cutlines using a blend buffer. + + Each output block is an independent pure function of its source windows and + the fixed source ordering, so blocks are processed in parallel. With + max_workers <= 1 (the default) processing is strictly serial and the output + is byte-for-byte identical to the original single-threaded loop. + + Args: + input_ortho_and_ortho_cuts: iterable of (orthophoto_path, cut_path) pairs. + output_orthophoto: path for the merged output GeoTIFF. + orthophoto_vars: rasterio profile overrides (TILED, COMPRESS, etc.). + merge_skip_blending: if True, skip the feather/cutline blend passes + (ODM #1934 --merge-skip-blending); only the first naive-copy pass runs. + max_workers: number of parallel worker threads (default 1 = serial). + Returns: + The output_orthophoto path, or None if there were no valid inputs. """ inputs = [] bounds=None @@ -293,8 +358,14 @@ def merge(input_ortho_and_ortho_cuts, output_orthophoto, orthophoto_vars={}, mer profile = first.profile num_bands = first.meta['count'] - 1 # minus alpha colorinterp = first.colorinterp + dst_count = first.count log.ODM_INFO("%s valid orthophoto rasters to merge" % len(inputs)) + + if max_workers > MERGE_WORKER_CAP: + log.ODM_INFO("Capping orthophoto merge workers %d -> %d to limit open file handles" % (max_workers, MERGE_WORKER_CAP)) + max_workers = MERGE_WORKER_CAP + sources = [(rasterio.open(o), rasterio.open(c)) for o,c in inputs] # scan input files. @@ -308,6 +379,10 @@ def merge(input_ortho_and_ortho_cuts, output_orthophoto, orthophoto_vars={}, mer if src.profile["count"] < 2: raise ValueError("Inputs must be at least 2-band rasters") dst_w, dst_s, dst_e, dst_n = min(xs), min(ys), max(xs), max(ys) + # Close the pre-scan handles; they are unused in the parallel block loop. + for s, c in sources: + s.close() + c.close() log.ODM_INFO("Output bounds: %r %r %r %r" % (dst_w, dst_s, dst_e, dst_n)) output_transform = Affine.translation(dst_w, dst_n) @@ -338,33 +413,73 @@ def merge(input_ortho_and_ortho_cuts, output_orthophoto, orthophoto_vars={}, mer if merge_skip_blending: log.ODM_INFO("Skipping second and third pass orthophoto blending, as --merge-skip-blending passed") - # create destination file - with rasterio.open(output_orthophoto, "w", **profile) as dstrast: + # create destination file. Cap GDAL's global block cache for the merge (restored on + # exit, see _bounded_gdal_cache): the default (5% of RAM) lets dirty output tiles + # accumulate, so a source read can trigger a large eviction/flush under GDAL's global + # lock that stalls the parallel readers; a small cache keeps flushes prompt and cheap. + cache_bytes = 512 * 1024 * 1024 + with _bounded_gdal_cache(cache_bytes), \ + rasterio.open(output_orthophoto, "w", **profile) as dstrast: dstrast.colorinterp = colorinterp - for idx, dst_window in dstrast.block_windows(): - left, bottom, right, top = dstrast.window_bounds(dst_window) + + # Each output block is an independent function of its source windows and + # the source ordering, so blocks can be COMPUTED in parallel. But writes + # to a single compressed, tiled GeoTIFF must happen in row-major (block) + # order: out-of-order writes cannot be flushed incrementally, so GDAL + # hoards every dirty block in RAM until it thrashes or OOMs. So we compute + # in a thread pool and write from one thread in strict block order, with a + # small bounded look-ahead for backpressure. max_workers <= 1 is a plain + # serial compute+write loop, identical to the original. + tls = threading.local() + block_windows = [(dst_window, dstrast.window_bounds(dst_window)) + for _, dst_window in dstrast.block_windows()] + total_blocks = len(block_windows) + log_every = max(1, total_blocks // 20) + + opened_sources = [] + opened_lock = threading.Lock() + + def get_sources(): + """Return this thread's (ortho, cut) rasterio dataset handles. + + GDAL/rasterio handles are not safe to share across threads, so each + worker thread lazily opens and caches its own set on first use and + registers it in opened_sources for cleanup after the parallel run. + """ + srcs = getattr(tls, "sources", None) + if srcs is None: + # In skip-blending mode the cut rasters are never read, so don't + # open them (halves the open file handles for this path). + srcs = [(rasterio.open(o), None if merge_skip_blending else rasterio.open(c)) + for o, c in inputs] + tls.sources = srcs + with opened_lock: + opened_sources.append(srcs) + return srcs + + def compute_block(item): + """Compute one output block (read + 3 blend passes); return the array. + + Does NOT write; writing happens in block order on the main thread. + """ + dst_window, (left, bottom, right, top) = item + local_sources = get_sources() blocksize = dst_window.width dst_rows, dst_cols = (dst_window.height, dst_window.width) - - # initialize array destined for the block - dst_count = first.count dst_shape = (dst_count, dst_rows, dst_cols) dstarr = np.zeros(dst_shape, dtype=dtype) # First pass, write all rasters naively without blending - for src, _ in sources: + for src, _ in local_sources: src_window = tuple(zip(rowcol( src.transform, left, top, op=round, precision=precision ), rowcol( src.transform, right, bottom, op=round, precision=precision ))) - temp = np.zeros(dst_shape, dtype=dtype) - temp = src.read( - out=temp, window=src_window, boundless=True, masked=False - ) + temp = _read_window_gated(src, src_window, dst_shape, dtype) # pixels without data yet are available to write write_region = np.logical_and( @@ -376,48 +491,41 @@ def merge(input_ortho_and_ortho_cuts, output_orthophoto, orthophoto_vars={}, mer if np.count_nonzero(dstarr[-1]) == blocksize: break - # Skip expensive blending operations if flag passed + # Skip the feather/cutline blend passes if requested (ODM #1934) if merge_skip_blending: - dstrast.write(dstarr, window=dst_window) - continue + return dstarr # Second pass, write all feathered rasters # blending the edges - for src, _ in sources: + for src, _ in local_sources: src_window = tuple(zip(rowcol( src.transform, left, top, op=round, precision=precision ), rowcol( src.transform, right, bottom, op=round, precision=precision ))) - temp = np.zeros(dst_shape, dtype=dtype) - temp = src.read( - out=temp, window=src_window, boundless=True, masked=False - ) + temp = _read_window_gated(src, src_window, dst_shape, dtype) where = temp[-1] != 0 for b in range(0, num_bands): blended = temp[-1] / 255.0 * temp[b] + (1 - temp[-1] / 255.0) * dstarr[b] np.copyto(dstarr[b], blended, casting='unsafe', where=where) dstarr[-1][where] = 255.0 - + # check if dest has any nodata pixels available if np.count_nonzero(dstarr[-1]) == blocksize: break # Third pass, write cut rasters # blending the cutlines - for _, cut in sources: + for _, cut in local_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 - ) + temp = _read_window_gated(cut, src_window, dst_shape, dtype) # For each band, average alpha values between # destination raster and cut raster @@ -425,6 +533,45 @@ def merge(input_ortho_and_ortho_cuts, output_orthophoto, orthophoto_vars={}, mer blended = temp[-1] / 255.0 * temp[b] + (1 - temp[-1] / 255.0) * dstarr[b] np.copyto(dstarr[b], blended, casting='unsafe', where=temp[-1]!=0) + return dstarr + + def write_block(idx, dst_window, dstarr): dstrast.write(dstarr, window=dst_window) + if (idx + 1) % log_every == 0: + log.ODM_INFO("Merging orthophoto: %s / %s blocks" % (idx + 1, total_blocks)) + + if max_workers <= 1: + # Serial: compute and write each block in order (original behavior). + for idx, item in enumerate(block_windows): + write_block(idx, item[0], compute_block(item)) + else: + # Parallel compute; one in-order writer with bounded look-ahead so at + # most `cap` blocks are in flight, which keeps memory small and gives GDAL + # strictly sequential, incrementally-flushable writes. + cap = max_workers * 2 + with ThreadPoolExecutor(max_workers=max_workers) as ex: + items = iter(enumerate(block_windows)) + pending = deque() + for _ in range(cap): + nxt = next(items, None) + if nxt is None: + break + idx, item = nxt + pending.append((idx, item[0], ex.submit(compute_block, item))) + while pending: + widx, wwin, fut = pending.popleft() + dstarr = fut.result() + write_block(widx, wwin, dstarr) + nxt = next(items, None) + if nxt is not None: + idx, item = nxt + pending.append((idx, item[0], ex.submit(compute_block, item))) + + # Close all thread-local source handles opened during the run. + for srcs in opened_sources: + for s, c in srcs: + s.close() + if c is not None: # cut handle is None in skip-blending mode + c.close() return output_orthophoto diff --git a/stages/splitmerge.py b/stages/splitmerge.py index 1b69c3f81..e808075d0 100644 --- a/stages/splitmerge.py +++ b/stages/splitmerge.py @@ -265,7 +265,7 @@ def process(self, args, outputs): os.remove(tree.odm_orthophoto_tif) orthophoto_vars = orthophoto.get_orthophoto_vars(args) - orthophoto.merge(all_orthos_and_ortho_cuts, tree.odm_orthophoto_tif, orthophoto_vars, args.merge_skip_blending) + orthophoto.merge(all_orthos_and_ortho_cuts, tree.odm_orthophoto_tif, orthophoto_vars, args.merge_skip_blending, max_workers=args.max_concurrency) orthophoto.post_orthophoto_steps(args, merged_bounds_file, tree.odm_orthophoto_tif, tree.orthophoto_tiles, args.orthophoto_resolution, reconstruction, tree, False) elif len(all_orthos_and_ortho_cuts) == 1: diff --git a/tests/test_orthophoto.py b/tests/test_orthophoto.py new file mode 100644 index 000000000..2792f7403 --- /dev/null +++ b/tests/test_orthophoto.py @@ -0,0 +1,103 @@ +import os +import shutil +import unittest + +import numpy as np +import rasterio +from rasterio.transform import Affine + +from opendm.orthophoto import _read_window_gated + + +class TestReadWindowGated(unittest.TestCase): + """_read_window_gated must be pixel-identical to a plain boundless read. + + It is an optimization that avoids rasterio's boundless VRT path for the + common in-bounds / fully-out-of-bounds cases, so for every window position + its result must equal what ds.read(..., boundless=True) would produce. + """ + + WIDTH = 10 + HEIGHT = 8 + COUNT = 4 # 3 bands + alpha, like the orthophotos merge() operates on + + def setUp(self): + self.tmp = "tests/assets/output" + if os.path.exists(self.tmp): + shutil.rmtree(self.tmp) + os.makedirs(self.tmp) + self.path = os.path.join(self.tmp, "gated_src.tif") + + # A distinct value per (band, row, col) so a wrong window/offset is visible. + self.data = np.arange( + self.COUNT * self.HEIGHT * self.WIDTH, dtype=np.uint8 + ).reshape(self.COUNT, self.HEIGHT, self.WIDTH) + + profile = { + "driver": "GTiff", + "width": self.WIDTH, + "height": self.HEIGHT, + "count": self.COUNT, + "dtype": "uint8", + "nodata": 0, + "transform": Affine.translation(100.0, 200.0) * Affine.scale(0.5, -0.5), + } + with rasterio.open(self.path, "w", **profile) as dst: + dst.write(self.data) + + def tearDown(self): + shutil.rmtree(self.tmp) + + def _boundless_reference(self, ds, window): + (r0, r1), (c0, c1) = window + shape = (self.COUNT, r1 - r0, c1 - c0) + out = np.zeros(shape, dtype=np.uint8) + return ds.read(out=out, window=window, boundless=True, masked=False) + + def _assert_matches_boundless(self, window): + (r0, r1), (c0, c1) = window + dst_shape = (self.COUNT, r1 - r0, c1 - c0) + with rasterio.open(self.path) as ds: + gated = _read_window_gated(ds, window, dst_shape, np.uint8) + reference = self._boundless_reference(ds, window) + self.assertEqual(gated.shape, dst_shape) + self.assertEqual(gated.dtype, np.uint8) + np.testing.assert_array_equal(gated, reference) + return gated + + def test_fully_inside(self): + gated = self._assert_matches_boundless(((2, 5), (3, 7))) + # Sanity: the fixture is non-trivial, so a real read is not all zeros. + self.assertTrue(gated.any()) + + def test_fully_inside_full_extent(self): + self._assert_matches_boundless(((0, self.HEIGHT), (0, self.WIDTH))) + + def test_fully_outside_left(self): + gated = self._assert_matches_boundless(((0, 3), (-5, -2))) + self.assertFalse(gated.any()) # nothing overlaps -> all zeros + + def test_fully_outside_right(self): + self._assert_matches_boundless(((0, 3), (self.WIDTH + 2, self.WIDTH + 5))) + + def test_fully_outside_top(self): + self._assert_matches_boundless(((-4, -1), (0, 3))) + + def test_fully_outside_bottom(self): + self._assert_matches_boundless(((self.HEIGHT + 1, self.HEIGHT + 4), (0, 3))) + + def test_partial_top_left_corner(self): + self._assert_matches_boundless(((-2, 3), (-2, 4))) + + def test_partial_bottom_right_corner(self): + self._assert_matches_boundless( + ((self.HEIGHT - 2, self.HEIGHT + 3), (self.WIDTH - 3, self.WIDTH + 2)) + ) + + def test_partial_straddles_full_width(self): + # In-bounds vertically, straddling both left and right edges at once. + self._assert_matches_boundless(((1, 4), (-2, self.WIDTH + 2))) + + +if __name__ == "__main__": + unittest.main()