Skip to content

pixano_inference.models.segmentation

Segmentation model base class.

The I/O types (:class:SegmentationInput/:class:SegmentationOutput) are the wire contract and live in :mod:pixano_inference.schemas.segmentation; they are re-exported here so from pixano_inference.models.segmentation import SegmentationInput keeps working.

SegmentationModel(config)

Bases: InferenceModel[SegmentationInput, SegmentationOutput]

Base class for image segmentation models.

predict receives a :class:SegmentationInput (image, prompts and options) and returns a :class:SegmentationOutput (masks, scores and, optionally, embeddings). The interface property, when overridden, returns a :class:SegmentationInterface: the prompts the model takes and what a prediction carries.

Example 
@register_model("my-segmenter")
class MySegmenter(SegmentationModel):
    def load_model(self):
        self.model = load_weights(self.config.model_params["path"])

    def predict(self, input: SegmentationInput) -> SegmentationOutput:
        masks, scores = self.model(input.image, input.points, input.labels, input.boxes)
        return SegmentationOutput(masks=masks, scores=scores)
Source code in pixano_inference/models/base.py 
53
54
55
56
57
58
59
def __init__(self, config: ModelDeploymentConfig) -> None:
    """Initialize the model with deployment config.

    Args:
        config: Model deployment configuration.
    """
    self._config = config