YOLOSeg wrapper

YOLOSEGBackend(arch, pretrained, num_classes=None, hyp=None)

Bases: torch.nn.Module

Ultralytics YOLO segmentation backend (v8/11-style SegmentationModel).

Builds the model from {arch}.yaml, optionally loads {arch}.pt weights, and stores loss hyperparameters (YoloSegLossParams) in self.model.args.

Parameters:

Name	Type	Description	Default
arch	str	Ultralytics segmentation architecture, e.g. "yolov8n-seg" or "yolo11m-seg".	required
pretrained	bool	If True, load Ultralytics pretrained weights through the public YOLO API.	required
num_classes	int | None	Number of classes; when None, read from current().data_instance_mask.num_classes.	None
hyp	deepvisiontools.models.yoloseg.yoloseg.YoloSegLossParams | None	Segmentation loss gains. If None, defaults are used.	None

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

def __init__(
    self,
    arch: str,
    pretrained: bool,
    num_classes: int | None = None,
    hyp: YoloSegLossParams | None = None,
) -> None:
    super().__init__()

    nc = int(
        num_classes
        if num_classes is not None
        else current().data_instance_mask.num_classes
    )

    self.model = SegmentationModel(f"{arch}.yaml", nc=nc)
    self.model.float()

    if pretrained:
        pretrained_model = _load_ultralytics_pretrained_model(arch)
        self.model.load(pretrained_model)

    cast(Any, self.model).args = hyp if hyp is not None else YoloSegLossParams()

    # Keep a runner object because build_yoloseg uses
    # backend._runner.model.init_criterion().
    runner = YOLO(f"{arch}.yaml")
    cast(Any, runner).model = self.model
    object.__setattr__(self, "_runner", runner)

    self.pad_mult = _stride_from_arch(arch)

forward(x, extra_args=None)

Single forward for train/eval.

Parameters:

Name	Type	Description	Default
x	Tensor	[B, 3, H, W] input batch.	required
extra_args	dict | None	Reserved for backend-specific extensions.	None

Returns:

Name	Type	Description
Any		Raw output structure consumed by the Ultralytics segmentation loss.

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

def forward(self, x: torch.Tensor, extra_args: None | dict = None):
    """
    Single forward for train/eval.

    Args:
        x (Tensor): `[B, 3, H, W]` input batch.
        extra_args (dict | None): Reserved for backend-specific extensions.

    Returns:
        Any: Raw output structure consumed by the Ultralytics segmentation loss.
    """
    return self.model(x)

YOLOSEGLoss(criterion, *, loss_factor=1.0)

Bases: deepvisiontools.models.base.basemodel.BaseLoss

Thin adapter for the Ultralytics YOLO-SEG loss.

Moves vendor criterion internals to the correct device. Aligns nested targets to the same device as predictions. Optionally scales reported losses by loss_factor.

Parameters:

Name	Type	Description	Default
criterion	typing.Any	Ultralytics seg loss object (e.g., from backend._runner.model.init_criterion()).	required
loss_factor	float	Overall divider to scale losses. Defaults to 1.0.	1.0

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

def __init__(self, criterion: Any, *, loss_factor: float = 1.0):
    super().__init__()
    self.criterion = criterion
    self.loss_factor = float(loss_factor)

to(*args, **kwargs)

Mirror nn.Module.to and also move internal vendor tensors to the target device.

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

def to(self, *args, **kwargs):
    """
    Mirror `nn.Module.to` and also move internal vendor tensors to the target device.
    """
    super().to(*args, **kwargs)
    device = kwargs.get("device")
    if device is None and len(args) >= 1:
        device = args[0]
    if device is not None:
        self._move_criterion_tensors(device)
    return self

forward(preds, targs)

Compute the segmentation loss and return a standardized LossDict.

Parameters:

Name	Type	Description	Default
preds	typing.Any	Raw vendor predictions (tensor or nested structure).	required
targs	Dict[str, Tensor]	Ultralytics-style target dict.	required

Returns:

Type	Description
dict[str, torch.Tensor]	Dict[str, Tensor]: Always contains "loss". When vendor returns
dict[str, torch.Tensor]	(total, detail), includes "loss_box", "loss_seg", "loss_cls",
dict[str, torch.Tensor]	and "loss_dfl" if available.

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

def forward(
    self, preds: Any, targs: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]:
    """
    Compute the segmentation loss and return a standardized `LossDict`.

    Args:
        preds (Any): Raw vendor predictions (tensor or nested structure).
        targs (Dict[str, Tensor]): Ultralytics-style target dict.

    Returns:
        Dict[str, Tensor]: Always contains `"loss"`. When vendor returns
        `(total, detail)`, includes `"loss_box"`, `"loss_seg"`, `"loss_cls"`,
        and `"loss_dfl"` if available.
    """
    device = self._first_tensor_device(preds)
    if device is not None:
        self._move_criterion_tensors(device)
        targs = self._align_targets_device(targs, device)

    batch_size = self._infer_batch_size(preds)
    factor = float(batch_size) * self.loss_factor

    out = self.criterion(preds, targs)

    if isinstance(out, tuple) and len(out) >= 2:
        raw_total = self._as_tensor(out[0], device)
        raw_detail = self._as_tensor(out[1], device)

        total = self._as_scalar_loss(raw_total) / factor
        detail = raw_detail / factor

        loss_dict: dict[str, torch.Tensor] = {"loss": total}

        try:
            detail_flat = detail.flatten()
            if detail_flat.numel() >= 4:
                loss_dict["loss_box"] = detail_flat[0]
                loss_dict["loss_seg"] = detail_flat[1]
                loss_dict["loss_cls"] = detail_flat[2]
                loss_dict["loss_dfl"] = detail_flat[3]
        except Exception:
            pass

        return loss_dict

    raw_loss = self._as_tensor(out, device)
    return {"loss": self._as_scalar_loss(raw_loss) / factor}

YoloSegLossParams(box=7.5, cls=0.5, dfl=1.5, overlap_mask=False) dataclass

Loss gain multipliers for Ultralytics v8/11 segmentation criterion. Note that in ultralytics, the segmentation gain used is the box gain.

Parameters:

Name	Type	Description	Default
box	float	IoU/box regression loss gain. Defaults to 7.5.	7.5
cls	float	Classification loss gain. Defaults to 0.5.	0.5
dfl	float	Distribution Focal Loss gain. Defaults to 1.5.	1.5
overlap_mask	bool	Whether to use overlapping mask loss (not supported by deepvisiontools yet). Defaults to False.	False

YOLOSEGPost()

Bases: torch.nn.Module

Decode YOLO-SEG raw outputs into BatchData[InstanceMaskData] (single-pass path).

Note

1) choose best class & confidence filter, 2) NMS in XYXY, 3) keep top-max_det, 4) project protos with coefficients to get mask logits; upsample to (H, W), 5) binarize at mask_logit_threshold, 6) pack into InstanceMaskData with labels/scores.

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

def __init__(self) -> None:
    super().__init__()
    self.conf_thresh: float = float(
        getattr(current(), "model_confidence_threshold", 0.25)
    )
    self.nms_iou: float = float(getattr(current(), "model_nms_iou", 0.7))
    self.max_det: int = int(getattr(current(), "model_max_detection", 300))
    self.mask_logit_threshold: float = float(
        getattr(current(), "model_mask_logit_threshold", 0.5)
    )

preds_from_raw(raw, aux=None)

Decode YOLO-seg outputs into BatchData[instance_mask].

Metric/eval path policy: - use permissive pred params supplied through aux["pred"]; - keep boxes clipped to the padded model canvas; - preserve separate instance masks in InstanceMaskData so metric accumulation does not lose overlapping-mask information.

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

def preds_from_raw(
    self, raw: tuple[torch.Tensor, ...], aux: dict | None = None
) -> BatchData:
    """
    Decode YOLO-seg outputs into `BatchData[instance_mask]`.

    Metric/eval path policy:
    - use permissive pred params supplied through `aux["pred"]`;
    - keep boxes clipped to the padded model canvas;
    - preserve separate instance masks in `InstanceMaskData` so metric
    accumulation does not lose overlapping-mask information.
    """
    if aux is None:
        aux = {}

    pp = cast(
        YoloSegPredParams | None,
        aux.get("pred", getattr(self, "default_pred_params", None)),
    )
    if pp is None:
        pp = YoloSegPredParams(
            conf_threshold=0.5,
            nms_iou=0.5,
            max_det=100,
            mask_logit_threshold=0.5,
        )

    canvas_size = aux.get("canvas_size")
    if canvas_size is None:
        raise ValueError("YOLOSEGPost requires aux['canvas_size'] for decoding.")

    H, W = int(canvas_size[0]), int(canvas_size[1])

    out0, out1 = self._unpack(raw)
    device = out0.device

    protos = out1[2]
    if not torch.is_tensor(protos):
        raise TypeError("YOLOSEGPost expected out1[2] to contain YOLO mask protos.")

    if out0.dim() != 3:
        raise TypeError(
            "YOLOSEGPost expected YOLO raw prediction tensor with shape "
            f"[B, D, N] or [B, N, D]. Got {tuple(out0.shape)}."
        )

    # Ultralytics usually returns [B, D, N], but some paths can expose [B, N, D].
    # D = 4 + num_classes + num_mask_protos.
    P = int(protos.shape[1])
    cfg_num_classes = int(
        getattr(current().data_instance_mask, "num_classes", 0) or 0
    )
    expected_d = 4 + cfg_num_classes + P if cfg_num_classes > 0 else None

    if expected_d is not None:
        if int(out0.shape[1]) == expected_d:
            out0 = out0.permute(0, 2, 1).contiguous()
        elif int(out0.shape[2]) == expected_d:
            out0 = out0.contiguous()
        elif int(out0.shape[1]) < int(out0.shape[2]):
            out0 = out0.permute(0, 2, 1).contiguous()
        else:
            out0 = out0.contiguous()
    else:
        if int(out0.shape[1]) < int(out0.shape[2]):
            out0 = out0.permute(0, 2, 1).contiguous()
        else:
            out0 = out0.contiguous()

    B, N, D = out0.shape
    C = int(D - 4 - P)
    if C <= 0:
        raise ValueError(
            f"Bad YOLO-seg head dimensions after layout normalization: "
            f"shape={tuple(out0.shape)}, P={P}, inferred C={C}."
        )

    boxes_cxcywh = out0[..., :4]  # [B, N, 4]
    cls_scores = out0[..., 4 : 4 + C]  # [B, N, C]
    coeffs = out0[..., 4 + C : 4 + C + P]  # [B, N, P]

    cls_scores = torch.nan_to_num(cls_scores, nan=-1e9)
    conf, labels = torch.max(cls_scores, dim=-1)  # [B, N], [B, N]

    boxes_f = boxes_cxcywh.reshape(-1, 4)
    coeffs_f = coeffs.reshape(-1, P)
    conf_f = conf.reshape(-1)
    labels_f = labels.reshape(-1).to(torch.long)

    image_ids = (
        torch.arange(B, device=device).unsqueeze(1).expand(B, N).reshape(-1)
    ).to(torch.long)

    # Permissive pre-filter. The metric suite will do the real confidence sweep.
    keep = conf_f >= float(pp.conf_threshold)
    boxes_f = boxes_f[keep]
    coeffs_f = coeffs_f[keep]
    conf_f = conf_f[keep]
    labels_f = labels_f[keep]
    image_ids = image_ids[keep]

    results: list[InstanceMaskData] = []

    if boxes_f.numel() == 0:
        for _ in range(B):
            results.append(InstanceMaskData.empty((H, W), device=device))
        return BatchData(results, device=device)

    # CXCYWH -> XYXY.
    cx, cy, bw, bh = boxes_f.unbind(dim=1)
    x1 = cx - 0.5 * bw
    y1 = cy - 0.5 * bh
    x2 = cx + 0.5 * bw
    y2 = cy + 0.5 * bh
    boxes_xyxy_f = torch.stack((x1, y1, x2, y2), dim=1).to(torch.float32)

    # Clip to the padded model canvas before NMS/mask projection.
    boxes_xyxy_f[:, 0] = boxes_xyxy_f[:, 0].clamp(0, W)
    boxes_xyxy_f[:, 2] = boxes_xyxy_f[:, 2].clamp(0, W)
    boxes_xyxy_f[:, 1] = boxes_xyxy_f[:, 1].clamp(0, H)
    boxes_xyxy_f[:, 3] = boxes_xyxy_f[:, 3].clamp(0, H)

    valid_box = (boxes_xyxy_f[:, 2] > boxes_xyxy_f[:, 0]) & (
        boxes_xyxy_f[:, 3] > boxes_xyxy_f[:, 1]
    )
    boxes_xyxy_f = boxes_xyxy_f[valid_box]
    coeffs_f = coeffs_f[valid_box]
    conf_f = conf_f[valid_box]
    labels_f = labels_f[valid_box]
    image_ids = image_ids[valid_box]

    if boxes_xyxy_f.numel() == 0:
        for _ in range(B):
            results.append(InstanceMaskData.empty((H, W), device=device))
        return BatchData(results, device=device)

    # Global class-aware batched NMS across the whole batch.
    # Fold image id into the class key so boxes from different images never
    # suppress each other.
    try:
        nms_ids = labels_f + image_ids * int(C)
        keep_idx = _batched_nms_amp_safe(
            boxes_xyxy_f,
            conf_f,
            nms_ids,
            float(pp.nms_iou),
        )
    except Exception:
        keep_chunks: list[torch.Tensor] = []

        for b in range(B):
            m_img = image_ids == b
            if not bool(m_img.any().item()):
                continue

            img_idx = torch.nonzero(m_img, as_tuple=False).squeeze(1)
            boxes_b = boxes_xyxy_f[m_img]
            conf_b = conf_f[m_img]
            labels_b = labels_f[m_img]

            for cls_id in labels_b.unique(sorted=True).tolist():
                m_cls = labels_b == int(cls_id)
                if not bool(m_cls.any().item()):
                    continue

                cls_idx_local = torch.nonzero(m_cls, as_tuple=False).squeeze(1)
                keep_local = _greedy_nms(
                    boxes_b[m_cls],
                    conf_b[m_cls],
                    float(pp.nms_iou),
                )
                keep_chunks.append(img_idx[cls_idx_local[keep_local]])

        if keep_chunks:
            keep_idx = torch.cat(keep_chunks, dim=0)
            keep_idx = keep_idx[torch.argsort(conf_f[keep_idx], descending=True)]
        else:
            keep_idx = torch.empty((0,), dtype=torch.long, device=device)

    boxes_xyxy_f = boxes_xyxy_f[keep_idx]
    coeffs_f = coeffs_f[keep_idx]
    conf_f = conf_f[keep_idx]
    labels_f = labels_f[keep_idx]
    image_ids = image_ids[keep_idx]

    # Global score ordering; per-image slices remain confidence-sorted.
    order = torch.argsort(conf_f, descending=True)
    boxes_xyxy_f = boxes_xyxy_f[order]
    coeffs_f = coeffs_f[order]
    conf_f = conf_f[order]
    labels_f = labels_f[order]
    image_ids = image_ids[order]

    for b in range(B):
        m = image_ids == b

        boxes_xyxy = boxes_xyxy_f[m]
        coeffs_b = coeffs_f[m]
        conf_b = conf_f[m]
        labels_b = labels_f[m]

        if boxes_xyxy.shape[0] > int(pp.max_det):
            max_det = int(pp.max_det)
            boxes_xyxy = boxes_xyxy[:max_det]
            coeffs_b = coeffs_b[:max_det]
            conf_b = conf_b[:max_det]
            labels_b = labels_b[:max_det]

        if boxes_xyxy.shape[0] == 0:
            results.append(InstanceMaskData.empty((H, W), device=device))
            continue

        proto = protos[b]  # [P, hp, wp]
        masks_prob = _proto2mask(proto, coeffs_b, boxes_xyxy, (H, W))
        masks_bin = masks_prob > float(pp.mask_logit_threshold)

        if masks_bin.shape[0] == 0:
            results.append(InstanceMaskData.empty((H, W), device=device))
            continue

        # Remove empty masks after projection/binarization.
        keep_nonempty = masks_bin.flatten(1).sum(1) > 0
        if not bool(keep_nonempty.any().item()):
            results.append(InstanceMaskData.empty((H, W), device=device))
            continue

        masks_bin = masks_bin[keep_nonempty]
        conf_b = conf_b[keep_nonempty]
        labels_b = labels_b[keep_nonempty]

        # Keep confidence ordering after empty-mask filtering.
        order_b = torch.argsort(conf_b, descending=True)
        masks_bin = masks_bin[order_b]
        conf_b = conf_b[order_b]
        labels_b = labels_b[order_b]

        # Important:
        # Preserve separate masks. InstanceMaskData will still build a stacked
        # value map for visualization, but metrics can consume the packed
        # separate-mask representation and avoid losing overlap information.
        imd = InstanceMaskData(
            value=None,
            labels=labels_b.to(torch.long),
            canvas_size=(H, W),
            scores=conf_b.to(torch.float32),
            separate_masks=masks_bin.to(torch.uint8),
            device=device,
        )
        imd = cast(InstanceMaskData, imd.sanitize())
        results.append(imd)

    return BatchData(results, device=device)

undo_preproc(preds, *, original_hw=None)

Undo input padding by center-cropping to the original size when provided.

Parameters:

Name	Type	Description	Default
preds	deepvisiontools.data.BatchData	Predictions after preds_from_raw.	required
original_hw	Tuple[int, int] | None	Original input (H, W). If None or there are no objects, returns preds unchanged.	None

Returns:

Name	Type	Description
BatchData	deepvisiontools.data.BatchData	Cropped predictions to (H, W) when needed.

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

def undo_preproc(
    self, preds: BatchData, *, original_hw: tuple[int, int] | None = None
) -> BatchData:
    """
    Undo input padding by center-cropping to the original size when provided.

    Args:
        preds (BatchData): Predictions after `preds_from_raw`.
        original_hw (Tuple[int, int] | None): Original input `(H, W)`. If `None` or
            there are no objects, returns `preds` unchanged.

    Returns:
        BatchData: Cropped predictions to `(H, W)` when needed.
    """
    if original_hw is None or preds.nb_object == 0:
        return preds
    orig_h, orig_w = original_hw
    canvas = preds.canvas_size
    if not canvas:
        return preds
    cur_h, cur_w = canvas
    if (orig_h, orig_w) == (cur_h, cur_w):
        return preds
    top = max((cur_h - orig_h) // 2, 0)
    left = max((cur_w - orig_w) // 2, 0)
    cropped = [d.crop(top, left, orig_h, orig_w) for d in preds.datas]
    return BatchData(cropped, device=preds.device)

YOLOSEGPre(pad_multiple)

Bases: torch.nn.Module

Preprocessing for YOLO-seg: - Pads images to the model stride multiple (e.g., 32/64/16). - Pads targets consistently and converts them to Ultralytics seg-loss format.

Parameters:

Name	Type	Description	Default
pad_multiple	int	Target stride multiple for padding.	required

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

def __init__(self, pad_multiple: int) -> None:
    super().__init__()
    self.pad_multiple = int(pad_multiple)

images_only(images)

Pad images so that height/width are divisible by pad_multiple.

Parameters:

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

Returns:

Name	Type	Description
Tensor	torch.Tensor	Padded images tensor.

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

def images_only(self, images: torch.Tensor) -> torch.Tensor:
    """
    Pad images so that height/width are divisible by `pad_multiple`.

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

    Returns:
        Tensor: Padded images tensor.
    """
    images, _ = _pad_to_multiple(images, self.pad_multiple)
    return images

images_targets(images, targets)

Pad images and targets, then produce Ultralytics-style segmentation targets.

Parameters:

Name	Type	Description	Default
images	Tensor	[B, C, H, W].	required
targets	deepvisiontools.data.BatchData	Instance-mask targets aligned with images.	required

Returns:

Type	Description
torch.Tensor	Tuple[Tensor, Dict[str, Tensor]]: Padded images and the target dict
dict[str, torch.Tensor]	{"batch_idx":[N,1], "cls":[N,1], "bboxes":[N,4], "masks":[B,H,W]} on the
tuple[torch.Tensor, dict[str, torch.Tensor]]	same device as images.

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

def images_targets(
    self, images: torch.Tensor, targets: BatchData
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """
    Pad images and targets, then produce Ultralytics-style segmentation targets.

    Args:
        images (Tensor): `[B, C, H, W]`.
        targets (BatchData): Instance-mask targets aligned with `images`.

    Returns:
        Tuple[Tensor, Dict[str, Tensor]]: Padded images and the target dict
        `{"batch_idx":[N,1], "cls":[N,1], "bboxes":[N,4], "masks":[B,H,W]}` on the
        **same device** as `images`.
    """
    images, (left, t, r, b) = _pad_to_multiple(images, self.pad_multiple)
    if targets.nb_object > 0:
        padded = [d.pad(left, t, r, b) for d in targets.datas]
        targets = BatchData(padded, device=targets.device)
    y = _targets_for_yoloseg(targets)
    # ensure correct device assigniation
    img_dev = images.device
    y = {k: (v.to(img_dev) if torch.is_tensor(v) else v) for k, v in y.items()}
    return images, y

YoloSegPredParams(conf_threshold, nms_iou, max_det, mask_logit_threshold) dataclass

Bases: deepvisiontools.models.base.basemodel.BasePredParams

Inference-time parameters for YOLO instance segmentation decoding.

Attributes:

Name	Type	Description
conf_threshold	float	Minimum class confidence required before NMS.
nms_iou	float	IoU threshold used by NMS (class-agnostic here).
max_det	int	Maximum number of instances to keep per image.
mask_logit_threshold	float	Threshold applied to mask logits to binarize instance masks (after projection from protos).

Notes

These parameters are passed by DeepVisionModel to YOLOSEGPost.preds_from_raw via aux={"pred": ...}, and are persisted automatically when saving the whole model object.

build_yoloseg(*, arch='yolov8n-seg', pretrained=True, num_classes=None, loss=None, loss_factor=1.0, hyp=None, supports_amp=True, autocast_dtype=torch.float16, device=None, use_amp=None, **extras)

Build a DeepVision YOLO instance segmentation wrapper, registered as "yoloseg".

This builder wraps an Ultralytics YOLO segmentation model inside the DeepVisionTools DeepVisionModel API. It wires YOLO-seg preprocessing, postprocessing, loss handling and prediction parameters so the model can be used with train_step, eval_step, predict, checkpointing and export helpers.

Parameters:

Name	Type	Description	Default
arch	Literal[...	YOLO segmentation architecture name. Supported values include YOLOv8-seg and YOLO11-seg variants, with optional -p6 variants. Defaults to "yolov8n-seg".	'yolov8n-seg'
pretrained	bool	If True, initialize the model from Ultralytics pretrained weights for the selected architecture. Defaults to True.	True
num_classes	int | None	Number of instance segmentation classes. When None, it is taken from current().data_instance_mask.num_classes. Defaults to None.	None
loss	typing.Any | None	Custom Ultralytics-compatible segmentation criterion. If None, the default vendor criterion is created from the internal YOLO runner. Defaults to None.	None
loss_factor	float	Divisor applied to the reported YOLO-seg losses. Useful to rescale vendor loss values. Defaults to 1.0.	1.0
hyp	deepvisiontools.models.yoloseg.yoloseg.YoloSegLossParams | None	YOLO-seg loss gains for box, cls, dfl and mask-related behavior. If None, default YoloSegLossParams values are used. Defaults to None.	None
supports_amp	bool	Whether this backend supports autocast mixed precision. Defaults to True.	True
autocast_dtype	torch.dtype	Autocast dtype used by DeepVisionModel. Defaults to torch.float16.	torch.float16
device	str | torch.device | None	Device used to place the model. If None, device resolution is delegated to DeepVisionModel. Defaults to None.	None
use_amp	bool | None	Force-enable or force-disable AMP. If None, the global/default DeepVisionTools AMP policy is used. Defaults to None.	None
extras	typing.Any	Additional keyword arguments kept for registry/export compatibility.	{}

Returns:

Name	Type	Description
DeepVisionModel	deepvisiontools.models.base.basemodel.DeepVisionModel	A fully wired instance segmentation model with
	deepvisiontools.models.base.basemodel.DeepVisionModel	instance_mask capabilities and default YoloSegPredParams loaded from the
	deepvisiontools.models.base.basemodel.DeepVisionModel	current configuration.

Export behavior

Saved artifacts are rebuilt with pretrained=False before loading the saved state dict. This avoids downloading vendor weights again when restoring a deepvisiontools checkpoint.

Usage

import torch

from deepvisiontools.models import ModelFactory

# Minimal YOLOv8 instance segmentation model
model = ModelFactory(
    name="yoloseg",
    arch="yolov8n-seg",
    pretrained=True,
    num_classes=3,
)

# YOLO11 instance segmentation model on CUDA
model = ModelFactory(
    name="yoloseg",
    arch="yolo11m-seg",
    pretrained=True,
    num_classes=3,
    device="cuda",
    use_amp=True,
)

# Larger stride variant
model = ModelFactory(
    name="yoloseg",
    arch="yolov8m-p6-seg",
    pretrained=True,
    num_classes=3,
)

# Custom loss gains
model = ModelFactory(
    name="yoloseg",
    arch="yolov8n-seg",
    num_classes=3,
    hyp=YoloSegLossParams(box=7.5, cls=0.5, dfl=1.5),
)

# Inference returns BatchData containing InstanceMaskData items.
images = torch.randn(2, 3, 512, 512)
preds = model.predict(images)
preds.data_type == "instance_mask"
>>> True

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

@MODEL_REGISTRY.register("yoloseg")
def build_yoloseg(
    *,
    arch: Literal[
        "yolov8n-seg",
        "yolov8m-seg",
        "yolov8l-seg",
        "yolov8x-seg",
        "yolov8n-p6-seg",
        "yolov8m-p6-seg",
        "yolov8l-p6-seg",
        "yolov8x-p6-seg",
        "yolo11n-seg",
        "yolo11m-seg",
        "yolo11l-seg",
        "yolo11x-seg",
    ] = "yolov8n-seg",
    pretrained: bool = True,
    num_classes: int | None = None,
    loss: Any | None = None,
    loss_factor: float = 1.0,
    hyp: YoloSegLossParams | None = None,
    supports_amp: bool = True,
    autocast_dtype: torch.dtype = torch.float16,
    device: str | torch.device | None = None,
    use_amp: bool | None = None,
    **extras: Any,
) -> DeepVisionModel:
    """
    Build a DeepVision YOLO instance segmentation wrapper, registered as `"yoloseg"`.

    This builder wraps an Ultralytics YOLO segmentation model inside the
    DeepVisionTools `DeepVisionModel` API. It wires YOLO-seg preprocessing,
    postprocessing, loss handling and prediction parameters so the model can be used
    with `train_step`, `eval_step`, `predict`, checkpointing and export helpers.

    Args:
        arch (Literal[..., optional): YOLO segmentation architecture name. Supported
            values include YOLOv8-seg and YOLO11-seg variants, with optional `-p6`
            variants. Defaults to `"yolov8n-seg"`.
        pretrained (bool, optional): If True, initialize the model from Ultralytics
            pretrained weights for the selected architecture. Defaults to True.
        num_classes (int | None, optional): Number of instance segmentation classes.
            When None, it is taken from `current().data_instance_mask.num_classes`.
            Defaults to None.
        loss (Any | None, optional): Custom Ultralytics-compatible segmentation
            criterion. If None, the default vendor criterion is created from the
            internal YOLO runner. Defaults to None.
        loss_factor (float, optional): Divisor applied to the reported YOLO-seg losses.
            Useful to rescale vendor loss values. Defaults to 1.0.
        hyp (YoloSegLossParams | None, optional): YOLO-seg loss gains for box, cls,
            dfl and mask-related behavior. If None, default `YoloSegLossParams`
            values are used. Defaults to None.
        supports_amp (bool, optional): Whether this backend supports autocast mixed
            precision. Defaults to True.
        autocast_dtype (torch.dtype, optional): Autocast dtype used by
            `DeepVisionModel`. Defaults to `torch.float16`.
        device (str | torch.device | None, optional): Device used to place the model.
            If None, device resolution is delegated to `DeepVisionModel`. Defaults to None.
        use_amp (bool | None, optional): Force-enable or force-disable AMP. If None,
            the global/default DeepVisionTools AMP policy is used. Defaults to None.
        extras: Additional keyword arguments kept for registry/export compatibility.

    Returns:
        DeepVisionModel: A fully wired instance segmentation model with
        `instance_mask` capabilities and default `YoloSegPredParams` loaded from the
        current configuration.

    ???+ note "Export behavior"
        Saved artifacts are rebuilt with `pretrained=False` before loading the saved
        state dict. This avoids downloading vendor weights again when restoring a
        deepvisiontools checkpoint.

    ???+ example "Usage"
        ```python
        import torch

        from deepvisiontools.models import ModelFactory

        # Minimal YOLOv8 instance segmentation model
        model = ModelFactory(
            name="yoloseg",
            arch="yolov8n-seg",
            pretrained=True,
            num_classes=3,
        )

        # YOLO11 instance segmentation model on CUDA
        model = ModelFactory(
            name="yoloseg",
            arch="yolo11m-seg",
            pretrained=True,
            num_classes=3,
            device="cuda",
            use_amp=True,
        )

        # Larger stride variant
        model = ModelFactory(
            name="yoloseg",
            arch="yolov8m-p6-seg",
            pretrained=True,
            num_classes=3,
        )

        # Custom loss gains
        model = ModelFactory(
            name="yoloseg",
            arch="yolov8n-seg",
            num_classes=3,
            hyp=YoloSegLossParams(box=7.5, cls=0.5, dfl=1.5),
        )

        # Inference returns BatchData containing InstanceMaskData items.
        images = torch.randn(2, 3, 512, 512)
        preds = model.predict(images)
        preds.data_type == "instance_mask"
        >>> True
        ```
    """

    # infer num classes once from config
    num_classes = (
        int(current().data_instance_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}.")

    backend = YOLOSEGBackend(
        arch=arch,
        pretrained=pretrained,
        num_classes=num_classes,
        hyp=hyp,
    )
    pre = YOLOSEGPre(pad_multiple=_stride_from_arch(arch))
    post = YOLOSEGPost()

    if loss is None:
        runner = cast(Any, backend._runner)
        criterion = runner.model.init_criterion()
        loss_fn = YOLOSEGLoss(criterion, loss_factor=loss_factor)
    else:
        loss_fn = YOLOSEGLoss(loss, loss_factor=loss_factor)

    caps = Capabilities(
        task="instance_mask",
        supports_amp=supports_amp,
        requires_divisible=_stride_from_arch(arch),
        description=f"Ultralytics YOLO-SEG:{arch}",
    )

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

    model.pred_params = YoloSegPredParams.default_from_config()

    # Artifact rebuild kwargs are deliberately side-effect-free. The checkpoint
    # state_dict will be loaded immediately after construction, so there is no need
    # and no benefit to re-loading Ultralytics pretrained weights here.
    export_kwargs = {
        "arch": arch,
        "num_classes": num_classes,
        **extras,
    }
    export_kwargs["pretrained"] = False

    model.set_export_spec("yoloseg", export_kwargs)

    return model