Skip to content

Reference for ultralytics/models/yolo/segment/predict.py

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/segment/predict.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!


ultralytics.models.yolo.segment.predict.SegmentationPredictor 

SegmentationPredictor(cfg=DEFAULT_CFG, overrides=None, _callbacks=None)

Bases: DetectionPredictor

A class extending the DetectionPredictor class for prediction based on a segmentation model.

This class specializes in processing segmentation model outputs, handling both bounding boxes and masks in the prediction results.

Attributes:

Name Type Description
args dict

Configuration arguments for the predictor.
model Module

The loaded YOLO segmentation model.
batch list

Current batch of images being processed.

Methods:

Name Description
postprocess

Applies non-max suppression and processes detections.
construct_results

Constructs a list of result objects from predictions.
construct_result

Constructs a single result object from a prediction.

Examples:

>>> from ultralytics.utils import ASSETS
>>> from ultralytics.models.yolo.segment import SegmentationPredictor
>>> args = dict(model="yolo11n-seg.pt", source=ASSETS)
>>> predictor = SegmentationPredictor(overrides=args)
>>> predictor.predict_cli()

This class specializes in processing segmentation model outputs, handling both bounding boxes and masks in the prediction results.

Parameters:

Name Type Description Default
cfg dict

Configuration for the predictor. Defaults to Ultralytics DEFAULT_CFG.

 DEFAULT_CFG
overrides dict

Configuration overrides that take precedence over cfg.

 None
_callbacks list

List of callback functions to be invoked during prediction.

 None
Source code in ultralytics/models/yolo/segment/predict.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize the SegmentationPredictor with configuration, overrides, and callbacks.

    This class specializes in processing segmentation model outputs, handling both bounding boxes and masks in the
    prediction results.

    Args:
        cfg (dict): Configuration for the predictor. Defaults to Ultralytics DEFAULT_CFG.
        overrides (dict, optional): Configuration overrides that take precedence over cfg.
        _callbacks (list, optional): List of callback functions to be invoked during prediction.
    """
    super().__init__(cfg, overrides, _callbacks)
    self.args.task = "segment"

construct_result 

construct_result(pred, img, orig_img, img_path, proto)

Construct a single result object from the prediction.

Parameters:

Name Type Description Default
pred ndarray

The predicted bounding boxes, scores, and masks.

 required
img Tensor

The image after preprocessing.

 required
orig_img ndarray

The original image before preprocessing.

 required
img_path str

The path to the original image.

 required
proto Tensor

The prototype masks.

 required

Returns:

Type Description
Results

Result object containing the original image, image path, class names, bounding boxes, and masks.
Source code in ultralytics/models/yolo/segment/predict.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def construct_result(self, pred, img, orig_img, img_path, proto):
    """
    Construct a single result object from the prediction.

    Args:
        pred (np.ndarray): The predicted bounding boxes, scores, and masks.
        img (torch.Tensor): The image after preprocessing.
        orig_img (np.ndarray): The original image before preprocessing.
        img_path (str): The path to the original image.
        proto (torch.Tensor): The prototype masks.

    Returns:
        (Results): Result object containing the original image, image path, class names, bounding boxes, and masks.
    """
    if not len(pred):  # save empty boxes
        masks = None
    elif self.args.retina_masks:
        pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
        masks = ops.process_mask_native(proto, pred[:, 6:], pred[:, :4], orig_img.shape[:2])  # HWC
    else:
        masks = ops.process_mask(proto, pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # HWC
        pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
    if masks is not None:
        keep = masks.sum((-2, -1)) > 0  # only keep predictions with masks
        pred, masks = pred[keep], masks[keep]
    return Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks)

construct_results 

construct_results(preds, img, orig_imgs, protos)

Construct a list of result objects from the predictions.

Parameters:

Name Type Description Default
preds List[Tensor]

List of predicted bounding boxes, scores, and masks.

 required
img Tensor

The image after preprocessing.

 required
orig_imgs List[ndarray]

List of original images before preprocessing.

 required
protos List[Tensor]

List of prototype masks.

 required

Returns:

Type Description
List[Results]

List of result objects containing the original images, image paths, class names, bounding boxes, and masks.
Source code in ultralytics/models/yolo/segment/predict.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def construct_results(self, preds, img, orig_imgs, protos):
    """
    Construct a list of result objects from the predictions.

    Args:
        preds (List[torch.Tensor]): List of predicted bounding boxes, scores, and masks.
        img (torch.Tensor): The image after preprocessing.
        orig_imgs (List[np.ndarray]): List of original images before preprocessing.
        protos (List[torch.Tensor]): List of prototype masks.

    Returns:
        (List[Results]): List of result objects containing the original images, image paths, class names,
            bounding boxes, and masks.
    """
    return [
        self.construct_result(pred, img, orig_img, img_path, proto)
        for pred, orig_img, img_path, proto in zip(preds, orig_imgs, self.batch[0], protos)
    ]

postprocess 

postprocess(preds, img, orig_imgs)

Apply non-max suppression and process segmentation detections for each image in the input batch.

Parameters:

Name Type Description Default
preds tuple

Model predictions, containing bounding boxes, scores, classes, and mask coefficients.

 required
img Tensor

Input image tensor in model format, with shape (B, C, H, W).

 required
orig_imgs list | Tensor | ndarray

Original image or batch of images.

 required

Returns:

Type Description
list

List of Results objects containing the segmentation predictions for each image in the batch. Each Results object includes both bounding boxes and segmentation masks.

Examples:

>>> predictor = SegmentationPredictor(overrides=dict(model="yolo11n-seg.pt"))
>>> results = predictor.postprocess(preds, img, orig_img)
Source code in ultralytics/models/yolo/segment/predict.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def postprocess(self, preds, img, orig_imgs):
    """
    Apply non-max suppression and process segmentation detections for each image in the input batch.

    Args:
        preds (tuple): Model predictions, containing bounding boxes, scores, classes, and mask coefficients.
        img (torch.Tensor): Input image tensor in model format, with shape (B, C, H, W).
        orig_imgs (list | torch.Tensor | np.ndarray): Original image or batch of images.

    Returns:
        (list): List of Results objects containing the segmentation predictions for each image in the batch.
               Each Results object includes both bounding boxes and segmentation masks.

    Examples:
        >>> predictor = SegmentationPredictor(overrides=dict(model="yolo11n-seg.pt"))
        >>> results = predictor.postprocess(preds, img, orig_img)
    """
    # Extract protos - tuple if PyTorch model or array if exported
    protos = preds[1][-1] if isinstance(preds[1], tuple) else preds[1]
    return super().postprocess(preds[0], img, orig_imgs, protos=protos)





📅 Created 1 year ago ✏️ Updated 7 months ago
jk4e glenn-jocher Burhan-Q