segrails.pipeline.inference

UNet inference pipeline step.

Loads a segmentation model from an MLflow model directory and runs inference on a raster PNG to produce a binary segmentation mask.

Classes

InferenceStep

Run UNet inference on a raster PNG and return a binary segmentation mask.

Module Contents

class segrails.pipeline.inference.InferenceStep(config: altametris.segrails.config.SegrailsConfig, model_dir: pathlib.Path, device: str = 'cpu')

Run UNet inference on a raster PNG and return a binary segmentation mask.

The model is loaded from the MLflow directory at construction time, moved to the target device and set to eval mode.

Parameters:

  • config – SegrailsConfig instance — provides encoder name, number of classes, inference image size, and normalisation statistics.

  • model_dir – Path to the MLflow model directory (MLmodel, data/, …) as returned by WeightStep.resolve_dir().

  • device – PyTorch device string. Defaults to "cpu". Pass "cuda" or "gpu" to run on GPU — falls back silently to CPU if CUDA is unavailable.

Example

>>> step = InferenceStep(config, model_dir=Path("/cache/U-Net-segrails"))
>>> mask, latency, image_np = step.predict(png_path)
_config
_device
model
predict(raster_path: pathlib.Path) → tuple[numpy.typing.NDArray[numpy.uint8], float, numpy.typing.NDArray[numpy.uint8]]

Run UNet inference on a raster PNG.

Parameters:

raster_path – Path to the input raster PNG produced by RasterStep.

Returns:

  • mask: uint8 numpy array of shape (H, W).

  • latency: inference duration in seconds.

  • image_np: RGB image as uint8 ndarray of shape (H, W, 3), resized to inference_image_size. Reusing this avoids a second file read in the post-processing step.

Return type:

A tuple of

Raises:

FileNotFoundError – If raster_path does not exist.

_load_model(model_dir: pathlib.Path) → altametris.unet.model.unet.Unet
_preprocess(raster_path: pathlib.Path) → torch.Tensor
_run_inference(image_tensor: torch.Tensor) → numpy.typing.NDArray[numpy.uint8]