Mapping XGBoost Predictions Back to a Raster Grid

TL;DR Answer

After you flatten a raster stack to a (n_pixels, n_bands) matrix and call model.predict(), you get a flat 1-D vector of predictions with no geographic meaning. To turn it back into a map you must: keep an integer index of the valid (non-nodata) pixels you fed the model, scatter the predictions into a full-length output array pre-filled with a nodata value, reshape that array to (rows, cols) using the same memory order you flattened with, and write a GeoTIFF that copies the source transform, crs, and dimensions. Getting the index and the reshape order consistent is the whole game — everything else is bookkeeping. This is the final step of the Gradient Boosting for Raster Data workflow.

import numpy as np
import rasterio

def scatter_back(preds, valid_idx, n_pixels, shape, nodata, dtype):
    out = np.full(n_pixels, nodata, dtype=dtype)
    out[valid_idx] = preds
    return out.reshape(shape)   # (rows, cols), C-order — matches the flatten

Part of: Gradient Boosting for Raster Data

Why the Round Trip Breaks

Training an XGBoost classifier on raster data means dissolving the grid into a table. A stack of n_bands aligned bands with rows × cols pixels becomes an (n_pixels, n_bands) array where every row is one pixel’s spectral signature. The model neither knows nor cares where each pixel sat — it sees rows, not geography. The moment you call predict, all spatial structure lives only in the order of the array. Lose the order and you lose the map.

Two things routinely corrupt that order. The first is nodata. Real rasters have masked pixels — cloud shadow, scene edges, water bodies clipped from the analysis. You cannot pass NaN or a sentinel value like -9999 into XGBoost and expect a meaningful class, so the standard pattern is to drop invalid pixels before prediction. That shrinks your prediction vector below n_pixels, and now a naive reshape fails outright or silently misaligns every subsequent pixel.

The second is memory order. NumPy flattens in C-order (row-major) by default, so array.reshape(-1) walks left-to-right, top-to-bottom. If any intermediate step reshapes in Fortran-order, or if you transpose a (bands, rows, cols) cube incorrectly, the predictions come back transposed — a map that looks like static. The fix is disciplined: flatten one way, keep an index, scatter back, reshape the same way. The rest of this page turns that discipline into a reusable, typed function.

Core Principles for a Reliable Scatter-Back

• Flatten spatially last. Read the raster as (bands, rows, cols), then move bands to the last axis and reshape to (rows*cols, bands). Record rows and cols immediately — they are the only way to rebuild the grid.
• Build the valid mask from every band. A pixel is usable only if it is valid in all bands. Combine per-band nodata masks with a logical AND so a cloud in band 3 also removes that pixel from bands 1 and 2.
• Keep the index, not just the mask. np.where(valid_mask)[0] gives you the integer positions of good pixels in the flat array. That index is what lets you scatter predictions back to their exact original slots.
• Pre-fill the output with nodata. Allocate the full n_pixels output first, filled with your nodata value, then write predictions only at the valid index. Dropped pixels stay nodata automatically.
• Match dtype to the task. Class labels belong in uint8 or uint16; probabilities and regression outputs belong in float32. Pick a nodata value the dtype can hold and that no real class uses.
• Copy the source profile, override deliberately. Inherit transform, crs, width, and height from the input. Change only count, dtype, nodata, and compression. Never hand-build a transform.

Production-Ready predict_to_raster()

The function below reads a source raster, masks invalid pixels, predicts with any fitted scikit-learn-style estimator (XGBoost included), scatters predictions back, and writes a compressed GeoTIFF that inherits the source georeferencing. It validates inputs, streams large rasters in windows, and never assumes the whole array fits in memory.

from __future__ import annotations

import logging
from pathlib import Path
from typing import Protocol

import numpy as np
import rasterio
from rasterio.windows import Window

logger = logging.getLogger(__name__)


class Estimator(Protocol):
    def predict(self, X: np.ndarray) -> np.ndarray: ...


def predict_to_raster(
    src_path: str | Path,
    dst_path: str | Path,
    model: Estimator,
    *,
    out_nodata: int | float = 255,
    out_dtype: str = "uint8",
    tile: int = 1024,
    compress: str = "deflate",
) -> Path:
    """Predict per-pixel with a fitted model and write a georeferenced GeoTIFF.

    The source raster is processed in square windows so arbitrarily large
    scenes never load into memory in full. Invalid (nodata) pixels are
    excluded from prediction and restored as out_nodata in the output.

    Parameters
    ----------
    src_path : str or Path
        Multi-band input raster. Bands become model features in order.
    dst_path : str or Path
        Output single-band class/probability map.
    model : Estimator
        Any object with a .predict(X) method (XGBClassifier, sklearn, ...).
    out_nodata : int or float, default=255
        Value written where the source pixel was nodata in any band.
    out_dtype : str, default="uint8"
        Output dtype. Use uint8/uint16 for class maps, float32 for scores.
    tile : int, default=1024
        Window edge length in pixels. Larger tiles = fewer, bigger batches.
    compress : str, default="deflate"
        Lossless GeoTIFF codec. Use deflate or lzw for class maps.

    Returns
    -------
    Path
        The path the raster was written to.
    """
    src_path, dst_path = Path(src_path), Path(dst_path)
    if not src_path.exists():
        raise FileNotFoundError(f"Source raster not found: {src_path}")
    if tile < 1:
        raise ValueError(f"tile must be a positive integer, got {tile}")

    with rasterio.open(src_path) as src:
        src_nodata = src.nodata
        profile = src.profile.copy()
        profile.update(
            count=1,
            dtype=out_dtype,
            nodata=out_nodata,
            compress=compress,
            tiled=True,
            blockxsize=256,
            blockysize=256,
            predictor=2,  # horizontal differencing helps class maps
        )

        with rasterio.open(dst_path, "w", **profile) as dst:
            for row_off in range(0, src.height, tile):
                for col_off in range(0, src.width, tile):
                    win = Window(
                        col_off,
                        row_off,
                        min(tile, src.width - col_off),
                        min(tile, src.height - row_off),
                    )
                    block = src.read(window=win)          # (bands, h, w)
                    out = _predict_block(
                        block, model, src_nodata, out_nodata, out_dtype
                    )
                    dst.write(out, 1, window=win)

    logger.info("Wrote prediction raster to %s", dst_path)
    return dst_path


def _predict_block(
    block: np.ndarray,
    model: Estimator,
    src_nodata: float | None,
    out_nodata: int | float,
    out_dtype: str,
) -> np.ndarray:
    """Flatten one window, predict valid pixels, scatter back to (h, w)."""
    n_bands, h, w = block.shape
    n_pixels = h * w

    # (bands, h, w) -> (h*w, bands), C-order so row/col order is preserved.
    flat = block.reshape(n_bands, n_pixels).transpose(1, 0)

    # A pixel is valid only if it is valid in every band.
    if src_nodata is None:
        valid_mask = np.ones(n_pixels, dtype=bool)
    else:
        valid_mask = np.all(flat != src_nodata, axis=1)
    valid_idx = np.where(valid_mask)[0]

    out = np.full(n_pixels, out_nodata, dtype=out_dtype)
    if valid_idx.size:
        preds = model.predict(flat[valid_idx])
        out[valid_idx] = preds.astype(out_dtype)

    return out.reshape(h, w)  # same C-order used to flatten

The key invariant is that reshape(n_bands, n_pixels) and the final reshape(h, w) both use C-order, and the transpose(1, 0) only swaps the band and pixel axes — never the spatial layout. That symmetry is what guarantees pixel (r, c) in equals pixel (r, c) out.

Step-by-Step Walkthrough

1. Fit or load your XGBoost model

Train exactly as you would for tabular data — the pixels are just rows. If your land-cover classes are unevenly represented, address that first with Handling Class Imbalance in Land Cover Classification, and tune the booster following Hyperparameter Tuning for XGBoost on Geodata.

import joblib
from xgboost import XGBClassifier

model: XGBClassifier = joblib.load("landcover_xgb.joblib")

2. Confirm the input bands match the training feature order

The model was trained on columns in a specific order. Band 1 of the raster must be feature 1, band 2 feature 2, and so on. A mismatch produces plausible-looking but wrong maps.

with rasterio.open("stack_2026.tif") as src:
    print("bands:", src.count, "| dtype:", src.dtypes[0])
    print("crs:", src.crs, "| nodata:", src.nodata)
    assert src.count == model.n_features_in_, "band/feature count mismatch"

3. Run the prediction to a GeoTIFF

out = predict_to_raster(
    "stack_2026.tif",
    "landcover_2026.tif",
    model,
    out_nodata=255,
    out_dtype="uint8",
    tile=2048,
)
print(f"Class map written to {out}")

4. Sanity-check the class distribution

A quick histogram catches gross failures — an all-one-class map, or a nodata value bleeding into real classes.

import numpy as np

with rasterio.open("landcover_2026.tif") as pred:
    arr = pred.read(1, masked=True)
    values, counts = np.unique(arr.compressed(), return_counts=True)
    for v, c in zip(values, counts):
        print(f"class {v}: {c:,} px ({100 * c / arr.count():.1f}%)")

If the source and prediction share a CRS with your other layers, you can overlay them directly. When they do not, reproject with Reprojecting a Raster with rasterio.warp.reproject rather than resampling the class labels with a continuous method.

How the Flat Vector Becomes a Map

The diagram traces one window through the pipeline: the band cube is flattened to a pixel table, invalid pixels are masked out and their positions recorded, the model predicts only the valid rows, and the results are scattered back into a nodata-filled grid before being written with the source georeferencing.

Verification

Never trust a prediction raster until you have proven it is aligned with its source. Open both files and assert that the georeferencing survived the round trip: identical shape, identical CRS, identical affine transform, and a nodata footprint that matches the input’s invalid pixels exactly.

import numpy as np
import rasterio


def verify_prediction(src_path: str, pred_path: str) -> None:
    """Assert the prediction raster inherited the source grid and georeferencing."""
    with rasterio.open(src_path) as src, rasterio.open(pred_path) as pred:
        assert (pred.height, pred.width) == (src.height, src.width), (
            f"shape mismatch: {pred.shape} vs {src.shape}"
        )
        assert pred.crs == src.crs, f"CRS mismatch: {pred.crs} vs {src.crs}"
        assert pred.transform.almost_equals(src.transform), "transform drifted"

        # The nodata footprint of the output must cover the source's invalid pixels.
        src_band = src.read(1, masked=True)
        pred_band = pred.read(1, masked=True)
        src_invalid = src_band.mask
        pred_invalid = pred_band.mask
        leaked = np.logical_and(src_invalid, ~pred_invalid).sum()
        assert leaked == 0, f"{leaked} source-nodata pixels got a real class"

    print("OK: shape, CRS, transform, and nodata footprint all preserved.")


verify_prediction("stack_2026.tif", "landcover_2026.tif")

If transform.almost_equals fails, you almost certainly rebuilt the profile by hand instead of copying src.profile. If the nodata assertion fails, your valid mask did not combine every band — a pixel valid in one band but nodata in another slipped through and received a fabricated class. Once this passes, the class map is safe to publish or hand off to a serving layer such as a Containerizing Geospatial Inference Pipelines workflow.

FAQ

Why do my predictions come out spatially scrambled when I write them to a raster?

Scrambled output almost always means the row/column order was lost between flattening and reshaping. NumPy’s default C-order ravel and reshape are self-consistent, but only if you never sort, shuffle, or drop rows without keeping a positional index. When you mask out nodata pixels before prediction, you must scatter the predictions back to their original flat positions using the saved valid_idx, then reshape with the same order used to flatten. Mixing C-order and Fortran-order between the two steps produces a transposed, unrecognisable map.

How do I preserve nodata pixels that were dropped before prediction?

Allocate a full-length output array pre-filled with your chosen nodata value, then assign predictions only at the valid-pixel index positions you recorded when masking. Pixels that were never passed to the model keep the nodata fill, so the georeferenced footprint of the output matches the input exactly. Set the same nodata value in the rasterio profile so downstream tools and GIS software mask those pixels automatically.

What dtype and compression should I use for a class map?

Store integer class labels as uint8 when you have fewer than 255 classes, reserving one value (commonly 255 or 0) for nodata. uint8 keeps files small and is universally supported. Apply lossless DEFLATE or LZW compression with tiling and predictor=2, which typically shrinks a categorical class map by five to ten times because adjacent pixels share the same label. Never use JPEG or any lossy codec on a class map — it invents fractional label values that no longer map to valid classes.

Related

• Hyperparameter Tuning for XGBoost on Geodata
• Handling Class Imbalance in Land Cover Classification
• Reprojecting a Raster with rasterio.warp.reproject
• Containerizing Geospatial Inference Pipelines

Part of: Gradient Boosting for Raster Data — Training Geospatial Predictive Models in Python