Skip to content

Reference for ultralytics/models/fastsam/predict.py

Note

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


ultralytics.models.fastsam.predict.FastSAMPredictor 

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

Bases: SegmentationPredictor

FastSAMPredictor is specialized for fast SAM (Segment Anything Model) segmentation prediction tasks in Ultralytics YOLO framework.

This class extends the SegmentationPredictor, customizing the prediction pipeline specifically for fast SAM. It adjusts post-processing steps to incorporate mask prediction and non-max suppression while optimizing for single- class segmentation.

Source code in ultralytics/models/yolo/segment/predict.py 
23
24
25
26
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """Initializes the SegmentationPredictor with the provided configuration, overrides, and callbacks."""
    super().__init__(cfg, overrides, _callbacks)
    self.args.task = "segment"

postprocess 

postprocess(preds, img, orig_imgs)

Applies box postprocess for FastSAM predictions.

Source code in ultralytics/models/fastsam/predict.py 
20
21
22
23
24
25
26
27
28
29
30
31
def postprocess(self, preds, img, orig_imgs):
    """Applies box postprocess for FastSAM predictions."""
    results = super().postprocess(preds, img, orig_imgs)
    for result in results:
        full_box = torch.tensor(
            [0, 0, result.orig_shape[1], result.orig_shape[0]], device=preds[0].device, dtype=torch.float32
        )
        boxes = adjust_bboxes_to_image_border(result.boxes.xyxy, result.orig_shape)
        idx = torch.nonzero(box_iou(full_box[None], boxes) > 0.9).flatten()
        if idx.numel() != 0:
            result.boxes.xyxy[idx] = full_box
    return results





Created 2023-11-12, Updated 2024-07-21
Authors: glenn-jocher (6), Burhan-Q (1)