Spaces:
Running
Running
| """Stage 1: DL inference. RGB (+ optional DEM) -> probability raster (ONNX). | |
| Ported from FraXtex's `fraxtex_tools/inference.py`, which is canonical for this | |
| pipeline (see CLAUDE.md). The logic below is deliberately kept line-comparable with | |
| the plugin: same contract reading, same Hann-blended overlapping tiles, same DEM | |
| handling, same output convention (high value = fracture). | |
| Two deliberate differences, both because the Space knows things the plugin cannot: | |
| * No input-mode enum. The plugin lets the user load any .onnx, so it must offer | |
| auto/rgb/rgb_dem to resolve what the file needs. Here the model comes from a fixed | |
| registry, so the channel count from the graph is always authoritative. | |
| * Input size limits (`limits.py`), because this container is shared. | |
| """ | |
| import numpy as np | |
| import onnxruntime as ort | |
| import rasterio | |
| from .limits import check_size | |
| def load_session(model_path): | |
| """Open an ONNX model and read the input contract out of it. | |
| Returns `(session, input_name, contract)` where contract carries the channel | |
| count, patch size and whether the caller must scale to [0, 1]. | |
| """ | |
| session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"]) | |
| input_name = session.get_inputs()[0].name | |
| shape = session.get_inputs()[0].shape | |
| # Symbolic dims (e.g. 'channels') are strings, so model_ch stays None. | |
| model_ch = shape[1] if len(shape) == 4 and isinstance(shape[1], int) else None | |
| # Tile size: the graph when it pins one (512 for SAM 2, 256 for the RGB+DEM | |
| # pair), else the historical default for dynamic exports. | |
| dim_h, dim_w = (shape[2], shape[3]) if len(shape) == 4 else (None, None) | |
| if isinstance(dim_h, int) and isinstance(dim_w, int): | |
| ps, ps_pinned = dim_h, True | |
| square = dim_h == dim_w | |
| else: | |
| ps, ps_pinned, square = 256, False, True | |
| meta = session.get_modelmeta().custom_metadata_map or {} | |
| return session, input_name, { | |
| "channels": model_ch, | |
| "patch_size": ps, | |
| "patch_pinned": ps_pinned, | |
| "square": square, | |
| "meta": meta, | |
| } | |
| def _resolve_input_range(contract, use_dem, feedback): | |
| """Decide whether the caller scales to [0, 1], from what the model declares. | |
| Preprocessing is not visible in the graph, so the model declares it: some models | |
| normalize internally and want raw 0-255. The channel count is the fallback for | |
| exports made before the keys existed. | |
| """ | |
| meta = contract["meta"] | |
| declared = (meta.get("input.range") or "").strip() | |
| if not declared and meta.get("preprocessing.in_graph", "").strip().lower() == "true": | |
| declared = "0-255" | |
| if declared in ("0-1", "0-255"): | |
| feedback.pushInfo(f"Model declares input range {declared}.") | |
| return declared == "0-1", False | |
| if declared: | |
| feedback.pushWarning( | |
| f"Ignoring unrecognised declared input range '{declared}'; " | |
| "expected '0-1' or '0-255'." | |
| ) | |
| scale_255 = use_dem | |
| feedback.pushInfo( | |
| "Model declares no input range; assuming " | |
| + ("0-1, scaling by /255" if scale_255 else "0-255, raw values") | |
| + f" from its {4 if use_dem else 3}-channel input." | |
| ) | |
| return scale_255, True | |
| def _read_rgb(rgb_path): | |
| """Read an orthophoto to HWC uint8. rasterio handles both PNG and GeoTIFF.""" | |
| with rasterio.open(rgb_path) as src: | |
| img = src.read() | |
| if img.shape[0] >= 3: | |
| img = img[:3].transpose(1, 2, 0) | |
| elif img.shape[0] == 2: | |
| img = img[0:1].repeat(3, axis=0).transpose(1, 2, 0) | |
| else: | |
| img = img.repeat(3, axis=0).transpose(1, 2, 0) | |
| if img.dtype != np.uint8: | |
| img = (img * 255).astype(np.uint8) if img.max() <= 1.0 else img.astype(np.uint8) | |
| return img, src.profile.copy(), src.crs | |
| def run_inference(rgb_path, model_path, feedback, dem_path=None, overlap=0.25): | |
| """Run tiled ONNX inference and return `(pred_uint8, profile, info)`. | |
| `pred_uint8` follows the plugin's convention: **high values mean fracture**. | |
| """ | |
| session, input_name, contract = load_session(model_path) | |
| model_ch = contract["channels"] | |
| ps = contract["patch_size"] | |
| if model_ch is not None and model_ch not in (3, 4): | |
| feedback.reportError( | |
| f"This model declares {model_ch} input channels. Only 3 (RGB) and " | |
| "4 (RGB+DEM) are supported.", | |
| fatalError=True, | |
| ) | |
| if not contract["square"]: | |
| feedback.reportError( | |
| "This model expects non-square inputs. Only square patches are supported.", | |
| fatalError=True, | |
| ) | |
| has_dem = dem_path is not None | |
| # The registry always yields a declared channel count; the None branch exists so | |
| # a third-party model dropped into the registry still behaves like the plugin. | |
| use_dem = has_dem if model_ch is None else model_ch == 4 | |
| if use_dem and not has_dem: | |
| feedback.reportError( | |
| "This model expects 4 input channels (RGB+DEM) but no DEM was provided. " | |
| "Provide a DEM, or choose an RGB-only model.", | |
| fatalError=True, | |
| ) | |
| if has_dem and not use_dem: | |
| feedback.pushWarning( | |
| "Running RGB-only: the supplied DEM is not fed to the model. It still " | |
| "provides the output georeferencing." | |
| ) | |
| scale_255, range_guessed = _resolve_input_range(contract, use_dem, feedback) | |
| img, rgb_profile, rgb_crs = _read_rgb(rgb_path) | |
| check_size(img.shape[0], img.shape[1], feedback) | |
| # The DEM donates the output profile whenever it is supplied, since the RGB may | |
| # be a plain PNG with no CRS. | |
| dem_arr = None | |
| if has_dem: | |
| with rasterio.open(dem_path) as src: | |
| profile = src.profile.copy() | |
| if (src.height, src.width) != img.shape[:2]: | |
| feedback.reportError( | |
| f"RGB and DEM must have the same dimensions " | |
| f"(RGB {img.shape[1]}x{img.shape[0]}, DEM {src.width}x{src.height}).", | |
| fatalError=True, | |
| ) | |
| if use_dem: | |
| dem_arr = src.read(1).astype(np.float32) | |
| if use_dem: | |
| dem_lo, dem_hi = float(dem_arr.min()), float(dem_arr.max()) | |
| feedback.pushInfo(f"DEM value range: [{dem_lo:.3f}, {dem_hi:.3f}]") | |
| if dem_lo < 0.0 or dem_hi > 255.0: | |
| feedback.pushWarning( | |
| "DEM values fall outside 0-255. The DEM band is scaled exactly like " | |
| "the RGB bands, so a DEM in raw metres enters the model far outside " | |
| "the range it was trained on. Byte-scale the DEM the way it was " | |
| "scaled for training." | |
| ) | |
| else: | |
| # A PNG source profile cannot be written back, so pin the driver. | |
| profile = rgb_profile | |
| profile.update(driver="GTiff") | |
| if rgb_crs is None: | |
| feedback.pushWarning( | |
| "The RGB orthophoto has no CRS, so the output probability raster is " | |
| "not georeferenced and vectorization works in pixel coordinates." | |
| ) | |
| # Overlapping patches, Hann-blended: no tile seams. | |
| if use_dem: | |
| combined = np.concatenate((img[:, :, :3], np.expand_dims(dem_arr, 2)), axis=2) | |
| del img, dem_arr # combined owns its data; the sources are dead weight now | |
| else: | |
| combined = img[:, :, :3] # a view, so img stays alive either way | |
| h, w = combined.shape[:2] | |
| n_ch = combined.shape[2] | |
| # Origins clamp to the far edge, so only sub-patch images need padding. | |
| pad_h = max(0, ps - h) | |
| pad_w = max(0, ps - w) | |
| if pad_h or pad_w: | |
| padded = np.pad(combined, ((0, pad_h), (0, pad_w), (0, 0)), mode="symmetric") | |
| else: | |
| padded = combined | |
| ph, pw = padded.shape[:2] | |
| stride = max(1, int(round(ps * (1.0 - overlap)))) | |
| def origins(size): | |
| offs = list(range(0, size - ps + 1, stride)) | |
| if offs[-1] != size - ps: | |
| offs.append(size - ps) # flush against the far edge | |
| return offs | |
| ys, xs = origins(ph), origins(pw) | |
| # Taper: ~0 at the patch border, 1 at its centre. Min weight ~1.5e-4, so | |
| # dividing a single-coverage pixel by it is safe in float32. | |
| w1d = np.hanning(ps + 2)[1:-1].astype(np.float32) | |
| win = np.outer(w1d, w1d) | |
| # win is separable over a full ys x xs grid, so the total weight per pixel is | |
| # outer(row_cov, col_cov), so no second full-size accumulator. | |
| def coverage(offs, size): | |
| cov = np.zeros(size, dtype=np.float32) | |
| for o in offs: | |
| cov[o:o + ps] += w1d | |
| return cov | |
| row_cov = coverage(ys, ph) | |
| col_cov = coverage(xs, pw) | |
| pred_sum = np.zeros((ph, pw), dtype=np.float32) | |
| total = len(ys) * len(xs) | |
| done = 0 | |
| feedback.pushInfo( | |
| f"Tiling: {len(xs)} x {len(ys)} = {total} patches " | |
| f"({ps} px, {int(overlap * 100)}% overlap, stride {stride}, {n_ch} channels)" | |
| ) | |
| for y in ys: | |
| for x in xs: | |
| patch = padded[y:y + ps, x:x + ps, :] | |
| patch_t = patch.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32) | |
| if scale_255: | |
| patch_t /= 255.0 | |
| if done == 0: | |
| try: | |
| out = session.run(None, {input_name: patch_t}) | |
| except Exception as exc: | |
| hint = ( | |
| f"The model declares {model_ch} input channels and was given " | |
| "that many, so the channel count is not the problem." | |
| + ( | |
| f" The {ps}x{ps} patch size came from the model itself." | |
| if contract["patch_pinned"] else | |
| f" It declares no patch size, so {ps}x{ps} was assumed." | |
| ) | |
| ) | |
| feedback.reportError( | |
| f"The model rejected a {n_ch}-channel {ps}x{ps} input: {exc}\n{hint}", | |
| fatalError=True, | |
| ) | |
| else: | |
| out = session.run(None, {input_name: patch_t}) | |
| pred_sum[y:y + ps, x:x + ps] += out[0][0, 0] * win | |
| done += 1 | |
| feedback.setProgress(int(done / total * 100)) | |
| # In place; both profiles are > 0 since every pixel is covered. | |
| pred_sum /= row_cov[:, np.newaxis] | |
| pred_sum /= col_cov[np.newaxis, :] | |
| pred = pred_sum[:h, :w] | |
| # Logits would wrap modulo 256 in the cast below. | |
| # | |
| # Deviation from the canonical plugin: it tests `lo < 0.0 or hi > 1.0` exactly. | |
| # A sigmoid that saturates to 1.0 (normal, and what unet-rgb does on the example | |
| # tile) picks up ~1.2e-07 of float32 error from the Hann blend and the coverage | |
| # division, so the exact test fires on every run of a perfectly good model. The | |
| # tolerance is far above that error and far below any real logit. Reported | |
| # upstream; revert this once the plugin carries the same tolerance. | |
| RANGE_TOL = 1e-5 | |
| lo, hi = float(pred.min()), float(pred.max()) | |
| if lo < -RANGE_TOL or hi > 1.0 + RANGE_TOL: | |
| feedback.pushWarning( | |
| f"Model output range [{lo:.3f}, {hi:.3f}] falls outside 0-1 and was clipped. " | |
| "This expects a model that emits probabilities (e.g. with a final sigmoid), " | |
| "not raw logits." | |
| ) | |
| # Clipped unconditionally, not just in the warning branch. Real logits are what | |
| # this guards: 8.0 * 255 = 2040 wraps to 248 in the uint8 cast below. The | |
| # sub-tolerance overshoot is harmless by comparison (255.00003 truncates to 255) | |
| # but clipping it too keeps the output provably in range. | |
| pred = np.clip(pred, 0.0, 1.0) | |
| # A near-constant map is the fingerprint of feeding the wrong input range. | |
| # Measured on the four published models: wrong range spans 0.000-0.016, right | |
| # range spans 0.79-1.00, so 0.05 sits in the gap between them. | |
| if range_guessed and hi - lo < 0.05: | |
| feedback.pushWarning( | |
| f"Output is nearly constant (range [{lo:.4f}, {hi:.4f}]) and this model " | |
| "declares no input range, so it was inferred from the channel count. The " | |
| "model may expect the other range." | |
| ) | |
| pred_uint8 = (pred * 255).astype(np.uint8) | |
| profile.update(count=1, dtype=rasterio.uint8, nodata=None, compress="lzw") | |
| info = { | |
| "crs": profile.get("crs"), | |
| "transform": profile.get("transform"), | |
| "patches": total, | |
| "patch_size": ps, | |
| "channels": n_ch, | |
| "used_dem": use_dem, | |
| "output_range": (lo, hi), | |
| } | |
| feedback.pushInfo("Inference complete.") | |
| return pred_uint8, profile, info | |