SMP wrapper

SemanticPredParams(binary_logits_threshold) dataclass

Bases: deepvisiontools.models.base.basemodel.BasePredParams

Prediction-time knobs for semantic segmentation decoding.

Attributes:

Name	Type	Description
binary_logits_threshold	float	Threshold applied to logits for binary segmentation (num_classes == 1). Not used for multiclass.

SMPBackend(arch, **kwargs)

Bases: torch.nn.Module

Thin wrapper around an SMP architecture. Accepts either an SMP class name (str) or an already-constructed torch.nn.Module. Exposes a single forward(x) that returns raw logits [B, C, H, W].

Parameters:

Name	Type	Description	Default
arch	str | torch.nn.Module	Either the SMP architecture name (e.g., "Unet", "DeepLabV3", "FPN", …) or a pre-built module instance.	required
**kwargs	typing.Any	Keyword args forwarded to the SMP constructor when arch is a string (e.g., encoder_name, encoder_weights, in_channels, classes, …).	{}

Source code in src/deepvisiontools/models/smp/smp.py

def __init__(self, arch: str | torch.nn.Module, **kwargs: Any) -> None:
    super().__init__()
    if isinstance(arch, str):
        ArchCls = getattr(smp, arch)
        self.model = ArchCls(**kwargs)
    else:
        self.model = arch  # already constructed

forward(x, extra_args=None)

Forward once through the SMP model.

Parameters:

Name	Type	Description	Default
x	torch.Tensor	Input image batch of shape [B, C, H, W].	required

Returns:

Type	Description
torch.Tensor	torch.Tensor: Raw logits of shape [B, C, H, W].

Source code in src/deepvisiontools/models/smp/smp.py

def forward(self, x: torch.Tensor, extra_args: None | dict = None) -> torch.Tensor:
    """
    Forward once through the SMP model.

    Args:
        x (torch.Tensor): Input image batch of shape `[B, C, H, W]`.

    Returns:
        torch.Tensor: Raw logits of shape `[B, C, H, W]`.
    """
    return self.model(x)

SMPLoss(smp_loss)

Bases: deepvisiontools.models.base.basemodel.BaseLoss

Adapter for losses from segmentation_models_pytorch.losses.

It accepts an already-constructed SMP loss object (e.g., FocalLoss, DiceLoss) and validates that it operates on logits (from_logits=True) since the backend returns raw scores.

Parameters:

Name	Type	Description	Default
smp_loss	torch.nn.Module	Instantiated SMP loss module. Must be compatible with the target tensor shape/dtype produced by SMPPre.images_targets.	required

Usage

import segmentation_models_pytorch as smp
loss = SMPLoss(smp.losses.FocalLoss(mode="binary"))  # logits-compatible
logits = torch.randn(2, 1, 64, 64)
target = torch.randint(0, 2, (2, 1, 64, 64)).float()
out = loss(logits, target)
isinstance(out, dict) and "loss" in out
>>> True

Source code in src/deepvisiontools/models/smp/smp.py

def __init__(self, smp_loss: torch.nn.Module):
    super().__init__()
    # Ensure losses that expose the flag are configured for logits
    fl = getattr(smp_loss, "from_logits", None)
    if fl is not None and fl is not True:
        raise ValueError(
            "SMPLoss expects raw logits from the backend, "
            "but the provided loss has from_logits=False. "
            "Please construct the SMP loss with from_logits=True."
        )
    self.criterion = smp_loss

forward(raw_logits, targets)

Compute the wrapped SMP loss.

Parameters:

Name	Type	Description	Default
raw_logits	torch.Tensor	Logits [B, C, H, W] from the backend.	required
targets	Union[torch.Tensor, deepvisiontools.data.BatchData]	Semantic masks as tensors or BatchData. When BatchData, this method attempts to tensorify by calling as_mask_tensor() if available, otherwise stacks items.	required

Returns:

Type	Description
	Dict[str, torch.Tensor]: {"loss": <scalar tensor>}.

Source code in src/deepvisiontools/models/smp/smp.py

def forward(self, raw_logits: torch.Tensor, targets):
    """
    Compute the wrapped SMP loss.

    Args:
        raw_logits (torch.Tensor): Logits `[B, C, H, W]` from the backend.
        targets (Union[torch.Tensor, BatchData]): Semantic masks as tensors or
            `BatchData`. When `BatchData`, this method attempts to tensorify
            by calling `as_mask_tensor()` if available, otherwise stacks items.

    Returns:
        Dict[str, torch.Tensor]: `{"loss": <scalar tensor>}`.
    """
    if hasattr(targets, "data_type"):  # BatchData
        # canonical path - if you have a helper to tensorify, prefer it.
        if hasattr(targets, "as_mask_tensor"):
            targets = targets.as_mask_tensor(device=raw_logits.device)
        else:
            # naive: stacks integer masks from BatchData.datas
            tensors = []
            for d in targets.datas:
                t = (
                    d.tensor
                    if torch.is_tensor(getattr(d, "tensor", None))
                    else d.to_tensor()
                )
                tensors.append(t)
            targets = torch.stack(tensors, dim=0).to(raw_logits.device)

    return {"loss": self.criterion(raw_logits, targets)}

SMPPost(num_classes=None)

Bases: torch.nn.Module

SMP postprocessor that converts raw logits [B, C, H, W] into BatchData by wrapping each item in SemanticLogits, which can further materialize SemanticMaskData as needed.

Note

• Uses current().data_semantic_mask.num_classes to parameterize decoding.
• undo_preproc is a no-op here (SMP path assumes spatial parity), but kept for API consistency with other wrappers.

Source code in src/deepvisiontools/models/smp/smp.py

def __init__(self, num_classes: int | None = None) -> None:
    super().__init__()
    self.num_classes = (
        int(current().data_semantic_mask.num_classes)
        if num_classes is None
        else int(num_classes)
    )
    if self.num_classes < 1:
        raise ValueError(f"num_classes must be >= 1, got {self.num_classes}.")

preds_from_raw(raw_outputs, aux=None)

Decode SMP raw logits to BatchData (single-pass path).

Parameters:

Name	Type	Description	Default
raw_outputs	torch.Tensor	Logits [B, C, H, W] from the backend.	required
aux	typing.Any	Additional context (unused here; accepted for API compatibility).	None

Returns:

Name	Type	Description
BatchData	deepvisiontools.data.BatchData	Collection of per-image semantic predictions.

Source code in src/deepvisiontools/models/smp/smp.py

def preds_from_raw(
    self, raw_outputs: torch.Tensor, aux: Any | None = None
) -> BatchData:
    """
    Decode SMP raw logits to `BatchData` (single-pass path).

    Args:
        raw_outputs (torch.Tensor): Logits `[B, C, H, W]` from the backend.
        aux (Any, optional): Additional context (unused here; accepted for API
            compatibility).

    Returns:
        BatchData: Collection of per-image semantic predictions.
    """
    logits = raw_outputs  # (B,C,H,W)
    preds = []
    for i in range(logits.shape[0]):
        sl = SemanticLogits(logits[i], self.num_classes)
        preds.append(sl.toSemanticMaskData())
    dev = logits.device.type if logits.is_cuda else "cpu"
    return BatchData(preds, device=dev)

undo_preproc(preds, *, original_hw=None)

Reverse preprocessing effects (no-op for SMP baseline).

Parameters:

Name	Type	Description	Default
preds	deepvisiontools.data.BatchData	Predictions produced by preds_from_raw.	required
original_hw	Tuple[int, int]	Original (H, W) image size if needed for resizing (ignored here).	None

Returns:

Name	Type	Description
BatchData	deepvisiontools.data.BatchData	Unmodified preds.

Source code in src/deepvisiontools/models/smp/smp.py

def undo_preproc(
    self, preds: BatchData, *, original_hw: tuple[int, int] | None = None
) -> BatchData:
    """
    Reverse preprocessing effects (no-op for SMP baseline).

    Args:
        preds (BatchData): Predictions produced by `preds_from_raw`.
        original_hw (Tuple[int, int], optional): Original `(H, W)` image size if
            needed for resizing (ignored here).

    Returns:
        BatchData: Unmodified `preds`.
    """
    return preds

SMPPre(num_classes=None)

Bases: torch.nn.Module

Prepares images and targets for SMP workflows.

• images_only is identity (expects [B, C, H, W]).
• images_targets adapts BatchData of semantic masks into tensors expected by SMP:
  • binary (num_classes == 1): targets → float tensor [B, 1, H, W];
  • multiclass: targets → long tensor [B, H, W].

The number of classes is read from current().data_semantic_mask.num_classes.

Source code in src/deepvisiontools/models/smp/smp.py

def __init__(self, num_classes: int | None = None) -> None:
    super().__init__()
    self.num_classes = (
        int(current().data_semantic_mask.num_classes)
        if num_classes is None
        else int(num_classes)
    )
    if self.num_classes < 1:
        raise ValueError(f"num_classes must be >= 1, got {self.num_classes}.")

images_only(images)

Identity preprocessing for images.

Parameters:

Name	Type	Description	Default
images	torch.Tensor	[B, C, H, W].	required

Returns:

Type	Description
torch.Tensor	torch.Tensor: Same tensor passed through unchanged.

Source code in src/deepvisiontools/models/smp/smp.py

def images_only(self, images: torch.Tensor) -> torch.Tensor:
    """
    Identity preprocessing for images.

    Args:
        images (torch.Tensor): `[B, C, H, W]`.

    Returns:
        torch.Tensor: Same tensor passed through unchanged.
    """
    return images

images_targets(images, targets)

Prepare (images, targets) for training/eval.

Parameters:

Name	Type	Description	Default
images	torch.Tensor	[B, C, H, W].	required
targets	deepvisiontools.data.BatchData	Semantic mask batch data.	required

Returns:

Type	Description
tuple[torch.Tensor, torch.Tensor]	Tuple[torch.Tensor, torch.Tensor]: (images, y) where: - for binary tasks: y is [B, 1, H, W] (float), - for multiclass: y is [B, H, W] (long).

Source code in src/deepvisiontools/models/smp/smp.py

def images_targets(
    self, images: torch.Tensor, targets: BatchData
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Prepare `(images, targets)` for training/eval.

    Args:
        images (torch.Tensor): `[B, C, H, W]`.
        targets (BatchData): Semantic mask batch data.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: `(images, y)` where:
            - for binary tasks: `y` is `[B, 1, H, W]` (float),
            - for multiclass:   `y` is `[B, H, W]` (long).
    """
    B, _, H, W = images.shape
    if targets.nb_object == 0:
        if self.num_classes == 1:
            y = torch.zeros((B, 1, H, W), dtype=torch.float32, device=images.device)
        else:
            y = torch.zeros((B, H, W), dtype=torch.long, device=images.device)
        return images, y

    masks: list[torch.Tensor] = [
        d.to(device=images.device) for d in targets.values()
    ]
    y = torch.stack(masks, dim=0)  # (B,H,W)

    y = y.unsqueeze(1).float() if self.num_classes == 1 else y.long()
    return images, y

build_smp(*, arch='Unet', encoder_name=None, encoder_weights=None, in_channels=3, num_classes=None, aux_params=None, loss=None, supports_amp=True, autocast_dtype=torch.float16, device=None, use_amp=None, **extra_arch_kwargs)

Build a DeepVision SMP wrapper (registered as "smp" in MODEL_REGISTRY).

This constructs an SMP model (e.g., Unet, DeepLabV3, FPN, PSPNet, …), wires the standard DeepVision pre/post/loss pipeline, and returns a DeepVisionModel ready for train_step, eval_step, and predict.

Parameters:

Name	Type	Description	Default
arch	str | torch.nn.Module	SMP architecture to wrap. Provide either the class name (e.g., "Unet", "DeepLabV3") or a pre-built module. Defaults to "Unet".	'Unet'
encoder_name	Optional[str]	Encoder backbone name (e.g., "resnet18"). See SMP docs for the full list. Defaults to None.	None
encoder_weights	Optional[str]	Pretrained weights name (e.g., "imagenet"). Defaults to None.	None
in_channels	int	Number of image channels. Defaults to 3.	3
num_classes	int | None	Number of output classes. When None, it is taken from current().data_semantic_mask.num_classes. Defaults to None.	None
aux_params	Optional[Dict[str, typing.Any]]	Additional SMP classifier head params (e.g., for auxiliary classification). Passed through to the SMP constructor. Defaults to None.	None
loss	Optional[torch.nn.Module]	SMP-compatible loss instance. If None, a default smp.losses.FocalLoss(mode=_infer_mode(num_classes)) is used. If the loss exposes .mode, it must match the inferred mode.	None
supports_amp	bool	Whether the backend supports autocast mixed precision. Defaults to True.	True
autocast_dtype	torch.dtype	Autocast dtype used by DeepVisionModel. Defaults to torch.float16.	torch.float16
**extra_arch_kwargs	typing.Any	Any additional keyword arguments forwarded to the SMP constructor (only when arch is a string).	{}

Returns:

Name	Type	Description
DeepVisionModel	deepvisiontools.models.base.basemodel.DeepVisionModel	Fully wired model (backend/pre/post/loss) with Capabilities
	deepvisiontools.models.base.basemodel.DeepVisionModel	populated and pred_params set via SemanticPredParams.default_from_config().

Example

# Minimal binary Unet
from deepvisiontools.models import ModelFactory
m = ModelFactory(
    name="smp",
    arch="Unet",
    encoder_name="resnet18",
    in_channels=3,
    num_classes=1,  # binary
)
# Multiclass with a custom loss
dice = smp.losses.DiceLoss(mode="multiclass", from_logits=True)
m2 = ModelFactory(
    name="smp",
    arch="FPN",
    encoder_name="resnet34",
    in_channels=3,
    num_classes=4,
    loss=dice,
)
# Already constructed SMP module
net = smp.Unet(encoder_name="resnet18", in_channels=1, classes=1)
m3 = ModelFactory(name="smp", arch=net)
# Training/eval/predict
x = torch.randn(2, 3, 128, 128)
y = torch.randint(0, 2, (2, 1, 128, 128)).float()
_ = m.train_step(x, y)  # doctest: +ELLIPSIS
_, preds = m.eval_step(x, y)  # doctest: +ELLIPSIS
preds.nb_object == x.shape[0]
>>> True

Source code in src/deepvisiontools/models/smp/smp.py

@MODEL_REGISTRY.register("smp")
def build_smp(
    *,
    arch: str | torch.nn.Module = "Unet",
    encoder_name: str | None = None,
    encoder_weights: str | None = None,
    in_channels: int = 3,
    num_classes: int | None = None,
    aux_params: dict[str, Any] | None = None,
    loss: Any | None = None,
    supports_amp: bool = True,
    autocast_dtype: torch.dtype = torch.float16,
    device: str | torch.device | None = None,
    use_amp: bool | None = None,
    **extra_arch_kwargs: Any,
) -> DeepVisionModel:
    """
    Build a DeepVision SMP wrapper (registered as `"smp"` in `MODEL_REGISTRY`).

    This constructs an SMP model (e.g., `Unet`, `DeepLabV3`, `FPN`, `PSPNet`, …),
    wires the standard DeepVision pre/post/loss pipeline, and returns a `DeepVisionModel`
    ready for `train_step`, `eval_step`, and `predict`.

    Args:
        arch (str | torch.nn.Module, optional): SMP architecture to wrap. Provide either
            the class name (e.g., `"Unet"`, `"DeepLabV3"`) or a pre-built module.
            Defaults to `"Unet"`.
        encoder_name (Optional[str], optional): Encoder backbone name (e.g., `"resnet18"`).
            See SMP docs for the full list. Defaults to `None`.
        encoder_weights (Optional[str], optional): Pretrained weights name (e.g., `"imagenet"`).
            Defaults to `None`.
        in_channels (int, optional): Number of image channels. Defaults to `3`.
        num_classes (int | None, optional): Number of output classes. When `None`, it is taken
            from `current().data_semantic_mask.num_classes`. Defaults to `None`.
        aux_params (Optional[Dict[str, Any]], optional): Additional SMP classifier head params
            (e.g., for auxiliary classification). Passed through to the SMP constructor.
            Defaults to `None`.
        loss (Optional[torch.nn.Module], optional): SMP-compatible loss instance. If `None`,
            a default `smp.losses.FocalLoss(mode=_infer_mode(num_classes))` is used.
            If the loss exposes `.mode`, it must match the inferred mode.
        supports_amp (bool, optional): Whether the backend supports autocast mixed precision.
            Defaults to `True`.
        autocast_dtype (torch.dtype, optional): Autocast dtype used by `DeepVisionModel`.
            Defaults to `torch.float16`.
        **extra_arch_kwargs: Any additional keyword arguments forwarded to the SMP constructor
            (only when `arch` is a string).

    Returns:
        DeepVisionModel: Fully wired model (backend/pre/post/loss) with `Capabilities`
        populated and `pred_params` set via `SemanticPredParams.default_from_config()`.

    ???+ example Usage
        ```Python
        # Minimal binary Unet
        from deepvisiontools.models import ModelFactory
        m = ModelFactory(
            name="smp",
            arch="Unet",
            encoder_name="resnet18",
            in_channels=3,
            num_classes=1,  # binary
        )
        # Multiclass with a custom loss
        dice = smp.losses.DiceLoss(mode="multiclass", from_logits=True)
        m2 = ModelFactory(
            name="smp",
            arch="FPN",
            encoder_name="resnet34",
            in_channels=3,
            num_classes=4,
            loss=dice,
        )
        # Already constructed SMP module
        net = smp.Unet(encoder_name="resnet18", in_channels=1, classes=1)
        m3 = ModelFactory(name="smp", arch=net)
        # Training/eval/predict
        x = torch.randn(2, 3, 128, 128)
        y = torch.randint(0, 2, (2, 1, 128, 128)).float()
        _ = m.train_step(x, y)  # doctest: +ELLIPSIS
        _, preds = m.eval_step(x, y)  # doctest: +ELLIPSIS
        preds.nb_object == x.shape[0]
        >>> True
        ```
    """
    classes_alias = extra_arch_kwargs.pop("classes", None)
    if num_classes is None and classes_alias is not None:
        num_classes = int(classes_alias)
    elif (
        num_classes is not None
        and classes_alias is not None
        and int(classes_alias) != int(num_classes)
    ):
        raise ValueError(
            "Conflicting class counts passed to build_smp: "
            f"num_classes={num_classes!r} and classes={classes_alias!r}. "
            "Use num_classes as the DeepVisionTools public API."
        )

    # num_classes default
    num_classes = (
        int(current().data_semantic_mask.num_classes)
        if num_classes is None
        else int(num_classes)
    )
    if num_classes < 1:
        raise ValueError(f"num_classes must be >= 1, got {num_classes}.")

    # Build SMP model kwargs
    arch_kwargs: dict[str, Any] = {
        "encoder_name": encoder_name,
        "encoder_weights": encoder_weights,
        "in_channels": in_channels,
        "classes": num_classes,
    }
    if aux_params is not None:
        arch_kwargs["aux_params"] = aux_params
    arch_kwargs.update({k: v for k, v in extra_arch_kwargs.items() if v is not None})

    backend = SMPBackend(arch=arch, **arch_kwargs)
    pre = SMPPre(num_classes=num_classes)
    post = SMPPost(num_classes=num_classes)

    # -----------------------
    # LOSS: runtime + portable spec
    # -----------------------
    inferred_mode = _infer_mode(num_classes)

    loss_obj: torch.nn.Module
    loss_spec: dict[str, Any]

    if loss is None:
        # Default loss (make it reproducible/portable)
        loss_obj = smp.losses.FocalLoss(mode=inferred_mode)
        loss_spec = _loss_to_spec(loss_obj)
    elif isinstance(loss, dict) and "target" in loss:
        # Portable form (used by checkpoint/yaml rebuild)
        loss_obj = _build_loss_from_spec(loss)
        loss_spec = {
            "target": str(loss["target"]),
            "kwargs": dict(loss.get("kwargs") or {}),
        }
    elif isinstance(loss, torch.nn.Module):
        # Runtime object (your current workflow)
        # Validate mode if exposed
        if hasattr(loss, "mode"):
            lm = loss.mode
            if lm != inferred_mode:
                raise ValueError(
                    f"Loss has inconsistent mode vs config: loss.mode={lm!r}, "
                    f"inferred={inferred_mode!r} from num_classes={num_classes}."
                )
        loss_obj = loss
        loss_spec = _loss_to_spec(loss_obj)
    else:
        raise TypeError(
            "build_smp(loss=...) expects None, a torch.nn.Module loss instance, "
            "or a portable dict {'target': 'module:QualName', 'kwargs': {...}}."
        )

    # SMPLoss wrapper keeps your current behavior (logits)
    loss_fn: SMPLoss = SMPLoss(loss_obj)

    caps = Capabilities(
        task="semantic_mask",
        supports_amp=supports_amp,
        requires_divisible=None,
        description=f"SMP:{arch}",
    )

    model = DeepVisionModel(
        backend=backend,
        pre=pre,
        post=post,
        loss_fn=loss_fn,
        caps=caps,
        device=device,
        use_amp=use_amp,
        autocast_dtype=autocast_dtype,
    )

    model.pred_params = SemanticPredParams.default_from_config()

    # IMPORTANT:
    # Export spec MUST be YAML-safe. So we store loss as a spec dict, not a Python object.
    model.set_export_spec(
        "smp",
        {
            "arch": arch,
            "encoder_name": encoder_name,
            # Artifact reload is immediately followed by state_dict loading.
            # Do not require SMP/timm to fetch encoder weights again.
            "encoder_weights": None,
            "in_channels": in_channels,
            "num_classes": num_classes,
            "aux_params": aux_params,
            "loss": loss_spec,
            **extra_arch_kwargs,
        },
    )

    return model