Add memory-efficient tiled inference with spline blending and D4 TTA

[jayakumarpujar]

Summary

Implements a generic, memory-efficient tiled inference pipeline for GeoTIFF rasters as a new geoai/inference.py module, addressing issue #87.

• Windowed rasterio I/O - reads tiles directly via rasterio windows, avoiding full-image input memory allocation
• Four blending strategies via BlendMode enum: none, linear, cosine, and spline (new, from issue author's approach using scipy.signal.windows.triang)
• D4 test-time augmentation - optional 8-fold augmentation (identity, 3 rotations, h-flip, v-flip, 2 diagonal flips) for improved predictions at 8x compute cost
• Configurable pre/post-processing hooks - preprocess_fn and postprocess_fn callables for model-specific normalization
• GeoTIFF metadata preservation - CRS, transform, and nodata are carried through to output

Key fixes vs prior attempt (PR #502)

• Dedicated geoai/inference.py module instead of adding to the already-large utils.py
• Proper boundary tile handling (crop predictions to actual tile dimensions, preventing shape mismatch)
• Input validation (overlap >= tile_size raises ValueError)
• num_classes parameter actually used to size the output accumulator
• No redundant device transfers
• Pure torch.rot90/torch.flip for D4 transforms (no torchvision dependency)

Files changed

File	Change
geoai/inference.py	New module: BlendMode, create_weight_mask, d4_forward, d4_inverse, d4_tta_forward, predict_geotiff
geoai/__init__.py	Lazy import registration for all public symbols
tests/test_inference.py	35 tests: imports, signatures, weight mask behavior, D4 roundtrips, validation
docs/inference.md	mkdocstrings API doc page
mkdocs.yml	Nav entry for inference module

Usage

from geoai.inference import predict_geotiff

predict_geotiff(
    model=my_model,
    input_raster="input.tif",
    output_raster="output.tif",
    tile_size=256,
    overlap=64,
    blend_mode="spline",
    tta=False,
)

Test plan

• All 35 unit tests pass (pytest tests/test_inference.py -v)
• Lazy imports verified (from geoai import predict_geotiff, BlendMode)
• Weight mask symmetry, shape, dtype, and value range validated for all 4 modes
• D4 roundtrip verified (forward + inverse = identity for all 8 transforms)
• Input validation tested (FileNotFoundError, ValueError for invalid overlap)

Closes #87

References

• Spline blending: https://github.com/Vooban/Smoothly-Blend-Image-Patches
• Original issue with source code: A more efficient inference function + smoothing #87

[giswqs]

@jayakumarpujar Thank you for adding this new feature. Can you add a notebook example to demostrate it?

[jayakumarpujar]

@giswqs I have updated with a notebook example. Please let me know if anything needs to be changed.

[Copilot]

Pull request overview

Adds a new geoai.inference module that provides a generic, windowed GeoTIFF inference pipeline with overlap blending and optional D4 test-time augmentation, plus docs/tests and top-level lazy exports.

Changes:

• Introduce geoai/inference.py with predict_geotiff, BlendMode, weight-mask generation, and D4 TTA helpers.
• Register new public API via lazy imports in geoai/__init__.py and add dedicated unit tests.
• Add mkdocs API page + example notebook and wire them into mkdocs.yml.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
geoai/inference.py	New tiled inference implementation (blending + D4 TTA) and public API symbols
geoai/__init__.py	Lazy-export registration for inference symbols and submodule
tests/test_inference.py	Unit tests for imports/exports, signatures, weight masks, D4 transforms, validation
docs/inference.md	mkdocstrings entry for the new inference module
docs/examples/smooth_inference.ipynb	Example notebook demonstrating blending strategies + TTA usage
mkdocs.yml	Navigation entries for the new notebook and inference docs page

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

_spline_window_1d() uses np.linspace(..., endpoint=False), so the last pixel in the taper never reaches weight 1.0, and the very next pixel jumps to 1.0 due to the center being initialized to ones. This introduces an avoidable discontinuity at the overlap boundary; consider using an endpoint-inclusive ramp (or explicitly setting the inner boundary weight to 1.0) so the taper joins the flat region continuously.

[jayakumarpujar]

Good catch! Fixed by switching to np.linspace(0.0, 1.0, overlap + 1, endpoint=True)[1:] so the taper reaches exactly 1.0 at the boundary - no more jump into the flat centre region. Added tests verifying continuity at the boundary for multiple overlap sizes.

[Copilot]

For BlendMode.SPLINE, _spline_window_1d() writes window[:overlap] and window[-overlap:] without guarding against overlap > window_size // 2. In that case the left/right taper regions overlap and overwrite each other, producing a distorted (and potentially non-monotone) window. Either constrain spline mode to overlap <= tile_size // 2 (raise a clear ValueError) or rework the construction to handle large overlaps similarly to the linear/cosine implementations.

[jayakumarpujar]

Valid point - the left/right tapers were overwriting each other when overlap > tile_size // 2. Added a ValueError guard for spline mode with a message suggesting linear or cosine blending for larger overlaps (those modes already handle this correctly via np.minimum).

[Copilot]

The code labels this as a “spline” window and references Smoothly-Blend-Image-Patches, but _spline_window_1d() is currently a powered raised-cosine taper (no triangular/spline windowing like the referenced approach). This mismatch is confusing for users trying to reproduce the referenced method; either implement the intended spline/triangular-based window or rename/re-document the mode to accurately describe the actual weighting function.

[jayakumarpujar]

Thanks for flagging. The naming is intentional - "spline" here refers to the smooth blending concept rather than the exact triangular window from the referenced repo. The docstring already described the actual implementation ("raised half-cosine shaped by power"). I've updated the enum docstring to say "powered raised-cosine taper" for additional clarity.

[Copilot]

output_dtype is configurable, but the default output_nodata=-9999.0 is not representable for many integer dtypes (e.g. uint8), which will likely cause rasterio to raise when opening the output dataset with profile.update(nodata=...). Consider validating that output_nodata fits output_dtype (and raising a clear error) or choosing a dtype-appropriate default nodata when output_dtype is changed.

[jayakumarpujar]

Great catch! Added validation that output_nodata fits within the valid range of output_dtype for integer types, raising a clear ValueError with guidance on valid alternatives.

[Copilot]

This notebook describes blend_mode="none" as “last-write-wins”, but the current predict_geotiff() implementation always normalizes accumulated predictions, so none results in uniform averaging across overlaps (not hard boundaries). Either update the notebook wording to match the actual behavior, or change the implementation so BlendMode.NONE truly overwrites in overlap regions.

[jayakumarpujar]

This is addressed together with comment 6 - the BlendMode.NONE docstring has been corrected to "uniform averaging", which is what the implementation actually does.

[Copilot]

BlendMode.NONE is documented as “No blending; last-write-wins”, but predict_geotiff() always does weighted accumulation + normalization, so blend_mode="none" actually averages overlapping tiles (uniform weights) rather than overwriting. Either implement true last-write-wins behavior for NONE (skip output_sum/weight_sum blending) or update the enum/docstrings (and docs notebook) to reflect that NONE means uniform averaging.

[jayakumarpujar]

You're right - the implementation uses weighted accumulation with uniform weights (all 1.0), which produces averaging in overlap regions, not last-write-wins. Updated the docstring to accurately say "Uniform averaging - all pixels are weighted equally (1.0), so overlapping tiles are simply averaged without tapering."