dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/models/sam/model.py

@@ -14,6 +14,8 @@ Key Features:
     - Trained on SA-1B dataset
 """
 
+from __future__ import annotations
+
 from pathlib import Path
 
 from ultralytics.engine.model import Model
@@ -23,12 +25,11 @@ from .predict import Predictor, SAM2Predictor
 
 
 class SAM(Model):
-    """
-    SAM (Segment Anything Model) interface class for real-time image segmentation tasks.
+    """SAM (Segment Anything Model) interface class for real-time image segmentation tasks.
 
-    This class provides an interface to the Segment Anything Model (SAM) from Ultralytics, designed for
-    promptable segmentation with versatility in image analysis. It supports various prompts such as bounding
-    boxes, points, or labels, and features zero-shot performance capabilities.
+    This class provides an interface to the Segment Anything Model (SAM) from Ultralytics, designed for promptable
+    segmentation with versatility in image analysis. It supports various prompts such as bounding boxes, points, or
+    labels, and features zero-shot performance capabilities.
 
     Attributes:
         model (torch.nn.Module): The loaded SAM model.
@@ -36,8 +37,8 @@ class SAM(Model):
         task (str): The task type, set to "segment" for SAM models.
 
     Methods:
-        predict: Performs segmentation prediction on the given image or video source.
-        info: Logs information about the SAM model.
+        predict: Perform segmentation prediction on the given image or video source.
+        info: Log information about the SAM model.
 
     Examples:
         >>> sam = SAM("sam_b.pt")
@@ -46,9 +47,8 @@ class SAM(Model):
         >>>     print(f"Detected {len(r.masks)} masks")
     """
 
-    def __init__(self, model="sam_b.pt") -> None:
-        """
-        Initialize the SAM (Segment Anything Model) instance.
+    def __init__(self, model: str = "sam_b.pt") -> None:
+        """Initialize the SAM (Segment Anything Model) instance.
 
         Args:
             model (str): Path to the pre-trained SAM model file. File should have a .pt or .pth extension.
@@ -66,8 +66,7 @@ class SAM(Model):
         super().__init__(model=model, task="segment")
 
     def _load(self, weights: str, task=None):
-        """
-        Load the specified weights into the SAM model.
+        """Load the specified weights into the SAM model.
 
         Args:
             weights (str): Path to the weights file. Should be a .pt or .pth file containing the model parameters.
@@ -81,17 +80,16 @@ class SAM(Model):
 
         self.model = build_sam(weights)
 
-    def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs):
-        """
-        Perform segmentation prediction on the given image or video source.
+    def predict(self, source, stream: bool = False, bboxes=None, points=None, labels=None, **kwargs):
+        """Perform segmentation prediction on the given image or video source.
 
         Args:
-            source (str | PIL.Image | numpy.ndarray): Path to the image or video file, or a PIL.Image object, or
-                a numpy.ndarray object.
+            source (str | PIL.Image | np.ndarray): Path to the image or video file, or a PIL.Image object, or a
+                np.ndarray object.
             stream (bool): If True, enables real-time streaming.
-            bboxes (List[List[float]] | None): List of bounding box coordinates for prompted segmentation.
-            points (List[List[float]] | None): List of points for prompted segmentation.
-            labels (List[int] | None): List of labels for prompted segmentation.
+            bboxes (list[list[float]] | None): List of bounding box coordinates for prompted segmentation.
+            points (list[list[float]] | None): List of points for prompted segmentation.
+            labels (list[int] | None): List of labels for prompted segmentation.
             **kwargs (Any): Additional keyword arguments for prediction.
 
         Returns:
@@ -108,20 +106,19 @@ class SAM(Model):
         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):
-        """
-        Perform segmentation prediction on the given image or video source.
+    def __call__(self, source=None, stream: bool = False, bboxes=None, points=None, labels=None, **kwargs):
+        """Perform segmentation prediction on the given image or video source.
 
-        This method is an alias for the 'predict' method, providing a convenient way to call the SAM model
-        for segmentation tasks.
+        This method is an alias for the 'predict' method, providing a convenient way to call the SAM model for
+        segmentation tasks.
 
         Args:
-            source (str | PIL.Image | numpy.ndarray | None): Path to the image or video file, or a PIL.Image
-                object, or a numpy.ndarray object.
+            source (str | PIL.Image | np.ndarray | None): Path to the image or video file, or a PIL.Image object, or a
+                np.ndarray object.
             stream (bool): If True, enables real-time streaming.
-            bboxes (List[List[float]] | None): List of bounding box coordinates for prompted segmentation.
-            points (List[List[float]] | None): List of points for prompted segmentation.
-            labels (List[int] | None): List of labels for prompted segmentation.
+            bboxes (list[list[float]] | None): List of bounding box coordinates for prompted segmentation.
+            points (list[list[float]] | None): List of points for prompted segmentation.
+            labels (list[int] | None): List of labels for prompted segmentation.
             **kwargs (Any): Additional keyword arguments to be passed to the predict method.
 
         Returns:
@@ -134,9 +131,8 @@ class SAM(Model):
         """
         return self.predict(source, stream, bboxes, points, labels, **kwargs)
 
-    def info(self, detailed=False, verbose=True):
-        """
-        Log information about the SAM model.
+    def info(self, detailed: bool = False, verbose: bool = True):
+        """Log information about the SAM model.
 
         Args:
             detailed (bool): If True, displays detailed information about the model layers and operations.
@@ -153,13 +149,12 @@ class SAM(Model):
         return model_info(self.model, detailed=detailed, verbose=verbose)
 
     @property
-    def task_map(self):
-        """
-        Provide a mapping from the 'segment' task to its corresponding 'Predictor'.
+    def task_map(self) -> dict[str, dict[str, type[Predictor]]]:
+        """Provide a mapping from the 'segment' task to its corresponding 'Predictor'.
 
         Returns:
-            (Dict[str, Dict[str, Type[Predictor]]]): A dictionary mapping the 'segment' task to its corresponding Predictor
-                class. For SAM2 models, it maps to SAM2Predictor, otherwise to the standard Predictor.
+            (dict[str, dict[str, Type[Predictor]]]): A dictionary mapping the 'segment' task to its corresponding
+                Predictor class. For SAM2 models, it maps to SAM2Predictor, otherwise to the standard Predictor.
 
         Examples:
             >>> sam = SAM("sam_b.pt")