Reference for ultralytics/models/sam/model.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/sam/model.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.models.sam.model.SAM
Bases: Model
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.
Attributes:
Name | Type | Description |
---|---|---|
model |
Module
|
The loaded SAM model. |
is_sam2 |
bool
|
Indicates whether the model is SAM2 variant. |
task |
str
|
The task type, set to "segment" for SAM models. |
Methods:
Name | Description |
---|---|
predict |
Performs segmentation prediction on the given image or video source. |
info |
Logs information about the SAM model. |
Examples:
>>> sam = SAM("sam_b.pt")
>>> results = sam.predict("image.jpg", points=[[500, 375]])
>>> for r in results:
>>> print(f"Detected {len(r.masks)} masks")
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model |
str
|
Path to the pre-trained SAM model file. File should have a .pt or .pth extension. |
'sam_b.pt'
|
Raises:
Type | Description |
---|---|
NotImplementedError
|
If the model file extension is not .pt or .pth. |
Examples:
Source code in ultralytics/models/sam/model.py
task_map
property
Provides a mapping from the 'segment' task to its corresponding 'Predictor'.
Returns:
Type | Description |
---|---|
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:
__call__
Performs 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.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
source |
str | Image | ndarray | None
|
Path to the image or video file, or a PIL.Image object, or a numpy.ndarray object. |
None
|
stream |
bool
|
If True, enables real-time streaming. |
False
|
bboxes |
List[List[float]] | None
|
List of bounding box coordinates for prompted segmentation. |
None
|
points |
List[List[float]] | None
|
List of points for prompted segmentation. |
None
|
labels |
List[int] | None
|
List of labels for prompted segmentation. |
None
|
**kwargs |
Any
|
Additional keyword arguments to be passed to the predict method. |
{}
|
Returns:
Type | Description |
---|---|
List
|
The model predictions, typically containing segmentation masks and other relevant information. |
Examples:
>>> sam = SAM("sam_b.pt")
>>> results = sam("image.jpg", points=[[500, 375]])
>>> print(f"Detected {len(results[0].masks)} masks")
Source code in ultralytics/models/sam/model.py
info
Logs information about the SAM model.
This method provides details about the Segment Anything Model (SAM), including its architecture, parameters, and computational requirements.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
detailed |
bool
|
If True, displays detailed information about the model layers and operations. |
False
|
verbose |
bool
|
If True, prints the information to the console. |
True
|
Returns:
Type | Description |
---|---|
Tuple
|
A tuple containing the model's information (string representations of the model). |
Examples:
Source code in ultralytics/models/sam/model.py
predict
Performs segmentation prediction on the given image or video source.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
source |
str | Image | ndarray
|
Path to the image or video file, or a PIL.Image object, or a numpy.ndarray object. |
required |
stream |
bool
|
If True, enables real-time streaming. |
False
|
bboxes |
List[List[float]] | None
|
List of bounding box coordinates for prompted segmentation. |
None
|
points |
List[List[float]] | None
|
List of points for prompted segmentation. |
None
|
labels |
List[int] | None
|
List of labels for prompted segmentation. |
None
|
**kwargs |
Any
|
Additional keyword arguments for prediction. |
{}
|
Returns:
Type | Description |
---|---|
List
|
The model predictions. |
Examples:
>>> sam = SAM("sam_b.pt")
>>> results = sam.predict("image.jpg", points=[[500, 375]])
>>> for r in results:
... print(f"Detected {len(r.masks)} masks")