Saltar al contenido

Referencia para ultralytics/models/sam/model.py

Nota

Este archivo est谩 disponible en https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/models/ sam/model .py. Si detectas alg煤n problema, por favor, ayuda a solucionarlo contribuyendo con una Pull Request 馃洜锔. 隆Gracias 馃檹!



ultralytics.models.sam.model.SAM

Bases: Model

SAM (Modelo de todo segmento).

SAM est谩 dise帽ado para la segmentaci贸n de im谩genes en tiempo real. Se puede utilizar con diversos indicadores, como cajas delimitadoras, puntos o etiquetas. El modelo puede funcionar sin disparos y se ha entrenado con el conjunto de datos SA-1B SA-1B.

C贸digo fuente en ultralytics/models/sam/model.py
class SAM(Model):
    """
    SAM (Segment Anything Model) interface class.

    SAM is designed for promptable real-time image segmentation. It can be used with a variety of prompts such as
    bounding boxes, points, or labels. The model has capabilities for zero-shot performance and is trained on the SA-1B
    dataset.
    """

    def __init__(self, model="sam_b.pt") -> None:
        """
        Initializes the SAM model with a pre-trained model file.

        Args:
            model (str): Path to the pre-trained SAM model file. File should have a .pt or .pth extension.

        Raises:
            NotImplementedError: If the model file extension is not .pt or .pth.
        """
        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):
        """
        Loads the specified weights into the SAM model.

        Args:
            weights (str): Path to the weights file.
            task (str, optional): Task name. Defaults to None.
        """
        self.model = build_sam(weights)

    def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs):
        """
        Performs segmentation prediction on the given image or video source.

        Args:
            source (str): Path to the image or video file, or a PIL.Image object, or a numpy.ndarray object.
            stream (bool, optional): If True, enables real-time streaming. Defaults to False.
            bboxes (list, optional): List of bounding box coordinates for prompted segmentation. Defaults to None.
            points (list, optional): List of points for prompted segmentation. Defaults to None.
            labels (list, optional): List of labels for prompted segmentation. Defaults to None.

        Returns:
            (list): The model predictions.
        """
        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):
        """
        Alias for the 'predict' method.

        Args:
            source (str): Path to the image or video file, or a PIL.Image object, or a numpy.ndarray object.
            stream (bool, optional): If True, enables real-time streaming. Defaults to False.
            bboxes (list, optional): List of bounding box coordinates for prompted segmentation. Defaults to None.
            points (list, optional): List of points for prompted segmentation. Defaults to None.
            labels (list, optional): List of labels for prompted segmentation. Defaults to None.

        Returns:
            (list): The model predictions.
        """
        return self.predict(source, stream, bboxes, points, labels, **kwargs)

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

        Args:
            detailed (bool, optional): If True, displays detailed information about the model. Defaults to False.
            verbose (bool, optional): If True, displays information on the console. Defaults to True.

        Returns:
            (tuple): A tuple containing the model's information.
        """
        return model_info(self.model, detailed=detailed, verbose=verbose)

    @property
    def task_map(self):
        """
        Provides a mapping from the 'segment' task to its corresponding 'Predictor'.

        Returns:
            (dict): A dictionary mapping the 'segment' task to its corresponding 'Predictor'.
        """
        return {"segment": {"predictor": Predictor}}

task_map property

Proporciona una correspondencia entre la tarea "Segmento" y su correspondiente "Predictor".

Devuelve:

Tipo Descripci贸n
dict

Un diccionario que asigna la tarea "Segmento" a su correspondiente "Predictor".

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

Alias del m茅todo "predecir".

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
source str

Ruta al archivo de imagen o v铆deo, o un objeto PIL.Image, o un objeto numpy.ndarray.

None
stream bool

Si es Verdadero, activa la transmisi贸n en tiempo real. Por defecto es Falso.

False
bboxes list

Lista de coordenadas de la caja delimitadora para la segmentaci贸n solicitada. Por defecto es Ninguno.

None
points list

Lista de puntos para la segmentaci贸n solicitada. Por defecto es Ninguno.

None
labels list

Lista de etiquetas para la segmentaci贸n solicitada. Por defecto es Ninguno.

None

Devuelve:

Tipo Descripci贸n
list

Las predicciones del modelo.

C贸digo fuente en ultralytics/models/sam/model.py
def __call__(self, source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs):
    """
    Alias for the 'predict' method.

    Args:
        source (str): Path to the image or video file, or a PIL.Image object, or a numpy.ndarray object.
        stream (bool, optional): If True, enables real-time streaming. Defaults to False.
        bboxes (list, optional): List of bounding box coordinates for prompted segmentation. Defaults to None.
        points (list, optional): List of points for prompted segmentation. Defaults to None.
        labels (list, optional): List of labels for prompted segmentation. Defaults to None.

    Returns:
        (list): The model predictions.
    """
    return self.predict(source, stream, bboxes, points, labels, **kwargs)

__init__(model='sam_b.pt')

Inicializa el modelo SAM con un archivo de modelo preentrenado.

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
model str

Ruta al archivo del modelo preentrenado SAM . El archivo debe tener una extensi贸n .pt o .pth.

'sam_b.pt'

Aumenta:

Tipo Descripci贸n
NotImplementedError

Si la extensi贸n del archivo del modelo no es .pt o .pth.

C贸digo fuente en ultralytics/models/sam/model.py
def __init__(self, model="sam_b.pt") -> None:
    """
    Initializes the SAM model with a pre-trained model file.

    Args:
        model (str): Path to the pre-trained SAM model file. File should have a .pt or .pth extension.

    Raises:
        NotImplementedError: If the model file extension is not .pt or .pth.
    """
    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")

info(detailed=False, verbose=True)

Registra informaci贸n sobre el modelo SAM .

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
detailed bool

Si es Verdadero, muestra informaci贸n detallada sobre el modelo. Por defecto es Falso.

False
verbose bool

Si es Verdadero, muestra informaci贸n en la consola. Por defecto es Verdadero.

True

Devuelve:

Tipo Descripci贸n
tuple

Una tupla que contiene la informaci贸n del modelo.

C贸digo fuente en ultralytics/models/sam/model.py
def info(self, detailed=False, verbose=True):
    """
    Logs information about the SAM model.

    Args:
        detailed (bool, optional): If True, displays detailed information about the model. Defaults to False.
        verbose (bool, optional): If True, displays information on the console. Defaults to True.

    Returns:
        (tuple): A tuple containing the model's information.
    """
    return model_info(self.model, detailed=detailed, verbose=verbose)

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

Realiza una predicci贸n de segmentaci贸n en la imagen o fuente de v铆deo dada.

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
source str

Ruta al archivo de imagen o v铆deo, o un objeto PIL.Image, o un objeto numpy.ndarray.

necesario
stream bool

Si es Verdadero, activa la transmisi贸n en tiempo real. Por defecto es Falso.

False
bboxes list

Lista de coordenadas de la caja delimitadora para la segmentaci贸n solicitada. Por defecto es Ninguno.

None
points list

Lista de puntos para la segmentaci贸n solicitada. Por defecto es Ninguno.

None
labels list

Lista de etiquetas para la segmentaci贸n solicitada. Por defecto es Ninguno.

None

Devuelve:

Tipo Descripci贸n
list

Las predicciones del modelo.

C贸digo fuente en ultralytics/models/sam/model.py
def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs):
    """
    Performs segmentation prediction on the given image or video source.

    Args:
        source (str): Path to the image or video file, or a PIL.Image object, or a numpy.ndarray object.
        stream (bool, optional): If True, enables real-time streaming. Defaults to False.
        bboxes (list, optional): List of bounding box coordinates for prompted segmentation. Defaults to None.
        points (list, optional): List of points for prompted segmentation. Defaults to None.
        labels (list, optional): List of labels for prompted segmentation. Defaults to None.

    Returns:
        (list): The model predictions.
    """
    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)





Creado 2023-11-12, Actualizado 2024-05-18
Autores: glenn-jocher (4), Burhan-Q (1)