Skip to content

Reference for ultralytics/models/sam/model.py

Note

Full source code for this file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/sam/model.py. Help us fix any issues you see by submitting a Pull Request 🛠️. Thank you 🙏!


ultralytics.models.sam.model.SAM

Bases: Model

SAM model interface.

Source code in ultralytics/models/sam/model.py
class SAM(Model):
    """
    SAM model interface.
    """

    def __init__(self, model='sam_b.pt') -> None:
        if model and Path(model).suffix not in ('.pt', '.pth'):
            raise NotImplementedError('SAM prediction requires pre-trained *.pt or *.pth model.')
        super().__init__(model=model, task='segment')

    def _load(self, weights: str, task=None):
        self.model = build_sam(weights)

    def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs):
        """Predicts and returns segmentation masks for given image or video source."""
        overrides = dict(conf=0.25, task='segment', mode='predict', imgsz=1024)
        kwargs.update(overrides)
        prompts = dict(bboxes=bboxes, points=points, labels=labels)
        return super().predict(source, stream, prompts=prompts, **kwargs)

    def __call__(self, source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs):
        """Calls the 'predict' function with given arguments to perform object detection."""
        return self.predict(source, stream, bboxes, points, labels, **kwargs)

    def info(self, detailed=False, verbose=True):
        """
        Logs model info.

        Args:
            detailed (bool): Show detailed information about model.
            verbose (bool): Controls verbosity.
        """
        return model_info(self.model, detailed=detailed, verbose=verbose)

    @property
    def task_map(self):
        return {'segment': {'predictor': Predictor}}

__call__(source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs)

Calls the 'predict' function with given arguments to perform object detection.

Source code in ultralytics/models/sam/model.py
def __call__(self, source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs):
    """Calls the 'predict' function with given arguments to perform object detection."""
    return self.predict(source, stream, bboxes, points, labels, **kwargs)

info(detailed=False, verbose=True)

Logs model info.

Parameters:

Name Type Description Default
detailed bool

Show detailed information about model.

False
verbose bool

Controls verbosity.

True
Source code in ultralytics/models/sam/model.py
def info(self, detailed=False, verbose=True):
    """
    Logs model info.

    Args:
        detailed (bool): Show detailed information about model.
        verbose (bool): Controls verbosity.
    """
    return model_info(self.model, detailed=detailed, verbose=verbose)

predict(source, stream=False, bboxes=None, points=None, labels=None, **kwargs)

Predicts and returns segmentation masks for given image or video source.

Source code in ultralytics/models/sam/model.py
def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs):
    """Predicts and returns segmentation masks for given image or video source."""
    overrides = dict(conf=0.25, task='segment', mode='predict', imgsz=1024)
    kwargs.update(overrides)
    prompts = dict(bboxes=bboxes, points=points, labels=labels)
    return super().predict(source, stream, prompts=prompts, **kwargs)




Created 2023-07-16, Updated 2023-08-07
Authors: glenn-jocher (5), Laughing-q (1)