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 |
|
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 |
|
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 |
|
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 |
|