segrails.pipeline.inference
===========================

.. py:module:: segrails.pipeline.inference

.. autoapi-nested-parse::

   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
-------

.. autoapisummary::

   segrails.pipeline.inference.InferenceStep


Module Contents
---------------

.. py:class:: 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.

   :param config: SegrailsConfig instance — provides encoder name, number of
                  classes, inference image size, and normalisation statistics.
   :param model_dir: Path to the MLflow model directory (``MLmodel``, ``data/``, …)
                     as returned by ``WeightStep.resolve_dir()``.
   :param device: PyTorch device string. Defaults to ``"cpu"``. Pass ``"cuda"``
                  or ``"gpu"`` to run on GPU — falls back silently to CPU if CUDA
                  is unavailable.

   .. rubric:: Example

   >>> step = InferenceStep(config, model_dir=Path("/cache/U-Net-segrails"))
   >>> mask, latency, image_np = step.predict(png_path)


   .. py:attribute:: _config


   .. py:attribute:: _device


   .. py:attribute:: model


   .. py:method:: predict(raster_path: pathlib.Path) -> tuple[numpy.typing.NDArray[numpy.uint8], float, numpy.typing.NDArray[numpy.uint8]]

      Run UNet inference on a raster PNG.

      :param 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.
      :rtype: A tuple of

      :raises FileNotFoundError: If ``raster_path`` does not exist.



   .. py:method:: _load_model(model_dir: pathlib.Path) -> altametris.unet.model.unet.Unet


   .. py:method:: _preprocess(raster_path: pathlib.Path) -> torch.Tensor


   .. py:method:: _run_inference(image_tensor: torch.Tensor) -> numpy.typing.NDArray[numpy.uint8]