Segmentation

optari.segmentation.segmenter

Segmentation model adapters.

ModelAdapterBase is the extension point for a new segmentation model: a subclass implements preprocess(), infer() and postprocess() for its own model format, and the model registry references it by class name in adapter_class.

ModelAdapterBase

ModelAdapterBase(model_config)

Bases: ABC

Base class for the segmentation adapters. Adapters have to implement preprocess, infer and postprocess methods, operating on single 2D grayscale frames.

Store the model config and its class id -> name mapping.

Source code in src/optari/segmentation/segmenter.py

def __init__(self, model_config: SegmentationModelConfig) -> None:
    """Store the model config and its class id -> name mapping."""
    self.model_config = model_config
    self.class_names = model_config.class_names

infer abstractmethod

infer(frame_2d)

Run the model on a preprocessed frame and return the raw predicted mask.

Source code in src/optari/segmentation/segmenter.py

@abstractmethod
def infer(self, frame_2d: np.ndarray) -> np.ndarray:
    """Run the model on a preprocessed frame and return the raw predicted mask."""
    raise NotImplementedError

postprocess abstractmethod

postprocess(mask_2d, frame_shape)

Turn the raw predicted mask into the final class id map for the original frame.

Implementations resize the mask back to frame_shape and clean up the class predictions.

Source code in src/optari/segmentation/segmenter.py

@abstractmethod
def postprocess(
    self, mask_2d: np.ndarray, frame_shape: tuple[int, int]
) -> np.ndarray:
    """Turn the raw predicted mask into the final class id map for the original frame.

    Implementations resize the mask back to *frame_shape* and clean up the class
    predictions.
    """
    raise NotImplementedError

predict

predict(frames)

Segment a batch of frames.

Parameters:

Name	Type	Description	Default
frames	ndarray	(nframes, H, W), use frames[np.newaxis] for a single frame.	required

Source code in src/optari/segmentation/segmenter.py

def predict(self, frames: np.ndarray) -> list[SegmentationResult]:
    """Segment a batch of frames.

    Args:
        frames: (nframes, H, W), use frames[np.newaxis] for a single frame.
    """
    results = []
    for frame_2d in frames:
        # A constant frame has nothing to segment, and its z-score normalization would divide
        # by a zero std and feed NaN to the model, which returns a meaningless mask.
        if np.ptp(frame_2d) == 0:
            results.append(
                SegmentationResult(
                    seg=np.zeros(frame_2d.shape, dtype=MASK_DTYPE),
                    class_names=dict(self.class_names),
                    blank=True,
                )
            )
        else:
            mask_2d = self.postprocess(
                self.infer(self.preprocess(frame_2d)), frame_2d.shape
            )
            results.append(
                SegmentationResult(
                    seg=mask_2d.astype(MASK_DTYPE),
                    class_names=dict(self.class_names),
                )
            )
    return results

preprocess abstractmethod

preprocess(frame_2d)

Prepare a single 2D grayscale frame for infer().

Implementations resize and normalize the frame to the model's expected input shape.

Source code in src/optari/segmentation/segmenter.py

@abstractmethod
def preprocess(self, frame_2d: np.ndarray) -> np.ndarray:
    """Prepare a single 2D grayscale frame for infer().

    Implementations resize and normalize the frame to the model's expected
    input shape.
    """
    raise NotImplementedError

SegmentationModelConfig dataclass

SegmentationModelConfig(model_id, filename, input_height, input_width, class_names, default_class, adapter_class, postprocessing_config=None, url=None)

Config for a segmentation model.

SegmentationResult dataclass

SegmentationResult(seg, class_names, blank=False)

Result of segmenting a single 2D US frame.

seg

(H, W) integer class id map.

class_names: Mapping of class id -> display name. blank: The frame was constant (e.g. a zero-padded frame that was never acquired) and got an empty mask instead of a prediction.

UKErUSSegAdapter

UKErUSSegAdapter(model_config)

Bases: ModelAdapterBase

Ultrasound segmentation adapter for scans taken at the university hospital of Erlangen using itheras Acuity Echo MSOT scanners. Requires trained ONNX model file and a config entry in segmentation_models.json.

Load the ONNX model weights and open an inference session.

Raises:

Type	Description
FileNotFoundError	The model weights have not been downloaded yet.

Source code in src/optari/segmentation/segmenter.py

def __init__(self, model_config: SegmentationModelConfig) -> None:
    """Load the ONNX model weights and open an inference session.

    Raises:
        FileNotFoundError: The model weights have not been downloaded yet.
    """
    super().__init__(model_config)

    self.input_shape = (
        model_config.input_height,
        model_config.input_width,
    )

    model_path = get_user_models_dir() / model_config.filename

    if not model_path.is_file():
        raise FileNotFoundError(
            f"Segmentation model not downloaded: {model_path}"
        )

    # CUDA is only available with onnxruntime-gpu (Linux). Elsewhere this resolves to CPU.
    providers = [
        p
        for p in ("CUDAExecutionProvider", "CPUExecutionProvider")
        if p in ort.get_available_providers()
    ]
    self.session = ort.InferenceSession(
        str(model_path), providers=providers
    )
    logger.info(
        "segmentation running on %s", self.session.get_providers()[0]
    )
    self.input_name = self.session.get_inputs()[0].name
    self.output_name = self.session.get_outputs()[0].name

infer

infer(frame_2d)

Run the ONNX session on the frame and return the per-pixel class id map.

Source code in src/optari/segmentation/segmenter.py

def infer(self, frame_2d: np.ndarray) -> np.ndarray:
    """Run the ONNX session on the frame and return the per-pixel class id map."""
    # add leading dim and run inference
    frame_tensor = np.expand_dims(frame_2d, axis=(0, 1))
    outputs = self.session.run(
        [self.output_name], {self.input_name: frame_tensor}
    )
    mask_2d = np.argmax(outputs[0], axis=1).squeeze()

    return mask_2d

postprocess

postprocess(mask_2d, frame_shape)

Perform custom postprocessing on the segmentation masks

Source code in src/optari/segmentation/segmenter.py

def postprocess(
    self, mask_2d: np.ndarray, frame_shape: tuple[int, int]
) -> np.ndarray:
    """
    Perform custom postprocessing on the segmentation masks
    """
    post_cfg = self.model_config.postprocessing_config
    # resize back to original frame shape
    mask_2d = resize_mask(mask_2d, frame_shape)
    # only keep the largest connected component for the given class ids
    mask_2d = keep_largest_region(
        mask_2d, post_cfg["keep_largest_per_class"]
    )
    # give freed pixels the most frequent remaining class in their row
    mask_2d = reassign_freed_pixels_row_based(mask_2d)
    # combine given classes into one class (here fascia classes)
    mask_2d = combine_classes(mask_2d, post_cfg["combine_class_groups"])
    for class_id, max_size in post_cfg["remove_small_objects_config"]:
        class_mask = mask_2d == class_id
        processed_class = remove_small_objects(
            class_mask, max_size=max_size
        )
        mask_2d[class_mask & ~processed_class] = 0
    # The row-wise step hands a row's majority class to pixels anywhere in that row,
    # which can recreate fragments of a class reduced to its largest component above.
    # Enforce that again last, then give the gaps their nearest class.
    mask_2d = keep_largest_region(
        mask_2d, post_cfg["keep_largest_per_class"]
    )
    mask_2d = reassign_freed_pixels(mask_2d)

    return mask_2d

preprocess

preprocess(frame_2d)

Resize the frame to the model's input shape and apply z-score normalization.

Source code in src/optari/segmentation/segmenter.py

def preprocess(self, frame_2d: np.ndarray) -> np.ndarray:
    """Resize the frame to the model's input shape and apply z-score normalization."""
    # resize to model input shape
    # TODO: check resize vs padding
    frame_2d = resize_img(frame_2d, self.input_shape)

    # z score normalization
    # TODO: check global mean/std vs per-frame
    frame_2d = normalize_img(frame_2d)

    return frame_2d

create_segmenter

create_segmenter(model_config)

Factory to create a segmentation adapter from a config.

Source code in src/optari/segmentation/segmenter.py

def create_segmenter(
    model_config: SegmentationModelConfig,
) -> ModelAdapterBase:
    """Factory to create a segmentation adapter from a config."""
    # Every adapter defined here is available by class name
    adapters = {cls.__name__: cls for cls in ModelAdapterBase.__subclasses__()}
    adapter_cls = adapters.get(model_config.adapter_class)
    if adapter_cls is None:
        raise ValueError(
            f"Unknown adapter class: {model_config.adapter_class}. "
            f"Available adapters: {sorted(adapters)}"
        )
    logger.info(
        "creating segmenter with model %s using adapter %s",
        model_config.model_id,
        model_config.adapter_class,
    )
    return adapter_cls(model_config)

load_model_registry

load_model_registry()

Load model configs and return a mapping of model_id -> config.

Source code in src/optari/segmentation/segmenter.py

def load_model_registry() -> dict[str, SegmentationModelConfig]:
    """Load model configs and return a mapping of model_id -> config."""

    model_config = json.loads(get_user_seg_models_config_file().read_text())
    models: dict[str, SegmentationModelConfig] = {}
    for item in model_config["models"]:
        model_id = str(item["id"])
        class_names = {int(k): str(v) for k, v in item["class_names"].items()}
        if max(class_names) > np.iinfo(MASK_DTYPE).max:
            raise ValueError(
                f"Segmentation model '{model_id}': class ids must be below "
                f"{np.iinfo(MASK_DTYPE).max + 1}"
            )
        models[model_id] = SegmentationModelConfig(
            model_id=model_id,
            filename=str(item["filename"]),
            class_names=class_names,
            input_height=int(item["input_height"]),
            input_width=int(item["input_width"]),
            default_class=str(item["default_class"]),
            adapter_class=str(item["adapter_class"]),
            postprocessing_config=item.get("postprocessing_config", None),
            url=item.get("url", None),
        )

    return models

optari.segmentation.segmentation_presets

Validation and normalization for segmentation preset settings.

validate_segmentation_settings

validate_segmentation_settings(settings, class_names=None)

Validate segmentation preset settings and normalize

Source code in src/optari/segmentation/segmentation_presets.py

def validate_segmentation_settings(
    settings: dict,
    class_names: Mapping[int, str] | None = None,
) -> dict:
    """Validate segmentation preset settings and normalize"""
    model_id = settings.get("model_id")
    if not isinstance(model_id, str) or not model_id.strip():
        raise ValueError("Segmentation preset requires a model_id.")

    selected_class_ids = settings.get("selected_class_ids")
    if not isinstance(selected_class_ids, list) or not selected_class_ids:
        raise ValueError("Segmentation preset requires selected class IDs.")
    if any(
        isinstance(class_id, bool) or not isinstance(class_id, int)
        for class_id in selected_class_ids
    ):
        raise ValueError("Selected class IDs must be integers.")
    if len(selected_class_ids) != len(set(selected_class_ids)):
        raise ValueError("Selected class IDs must be unique.")

    roi_class_id = settings.get("roi_class_id")
    if roi_class_id is not None and (
        isinstance(roi_class_id, bool) or not isinstance(roi_class_id, int)
    ):
        raise ValueError("ROI class ID must be an integer or null.")

    roi_shape = settings.get("roi_shape", "ellipse")
    if not isinstance(roi_shape, str) or roi_shape not in SUPPORTED_ROI_SHAPES:
        raise ValueError(f"Unsupported ROI shape: {roi_shape}")

    roi_width_mm = _optional_dimension(
        settings.get("roi_width_mm"), "ROI width"
    )
    roi_height_mm = _optional_dimension(
        settings.get("roi_height_mm"), "ROI height"
    )
    if (
        roi_class_id is not None
        and roi_shape in {"ellipse", "rectangle"}
        and (
            roi_width_mm is None
            or roi_height_mm is None
            or roi_width_mm <= 0
            or roi_height_mm <= 0
        )
    ):
        raise ValueError(
            "Ellipse and rectangle presets require positive width and height."
        )

    if class_names is not None:
        available_ids = set(class_names)
        unknown_ids = set(selected_class_ids) - available_ids
        if unknown_ids:
            raise ValueError(
                f"Preset contains unknown class IDs: {sorted(unknown_ids)}"
            )
        if roi_class_id is not None:
            if roi_class_id not in available_ids:
                raise ValueError(
                    f"Preset contains unknown ROI class ID: {roi_class_id}"
                )
            if roi_class_id not in selected_class_ids:
                raise ValueError(
                    "ROI class ID must be one of the selected classes."
                )

    return {
        "model_id": model_id,
        "selected_class_ids": list(selected_class_ids),
        "roi_class_id": roi_class_id,
        "roi_shape": roi_shape,
        "roi_width_mm": roi_width_mm,
        "roi_height_mm": roi_height_mm,
        "roi_top_margin_mm": _optional_dimension(
            settings.get("roi_top_margin_mm"), "ROI top margin"
        ),
    }