mirror of
https://github.com/ultralytics/ultralytics.git
synced 2025-09-15 15:48:41 +08:00
Add MobileSAM auto annotation feature 🚀 (#18374)
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
parent
d88e923595
commit
93862d3640
12
docs/en/macros/sam-auto-annotate.md
Normal file
12
docs/en/macros/sam-auto-annotate.md
Normal file
@ -0,0 +1,12 @@
|
||||
| Argument | Type | Default | Description |
|
||||
| ------------ | ----------- | -------------- | --------------------------------------------------------------------------------- |
|
||||
| `data` | `str` | required | Path to directory containing target images/videos for annotation or segmentation. |
|
||||
| `det_model` | `str` | `"yolo11x.pt"` | YOLO detection model path for initial object detection. |
|
||||
| `sam_model` | `str` | `"sam2_b.pt"` | SAM2 model path for segmentation (supports t/s/b/l variants and SAM2.1 models). |
|
||||
| `device` | `str` | `""` | Computation device (e.g., 'cuda:0', 'cpu', or '' for automatic device detection). |
|
||||
| `conf` | `float` | `0.25` | YOLO detection confidence threshold for filtering weak detections. |
|
||||
| `iou` | `float` | `0.45` | IoU threshold for Non-Maximum Suppression to filter overlapping boxes. |
|
||||
| `imgsz` | `int` | `640` | Input size for resizing images (must be multiple of 32). |
|
||||
| `max_det` | `int` | `300` | Maximum number of detections per image for memory efficiency. |
|
||||
| `classes` | `list[int]` | `None` | List of class indices to detect (e.g., `[0, 1]` for person & bicycle). |
|
||||
| `output_dir` | `str` | `None` | Save directory for annotations (defaults to './labels' relative to data path). |
|
||||
@ -130,6 +130,22 @@ You can download the model [here](https://github.com/ChaoningZhang/MobileSAM/blo
|
||||
|
||||
We have implemented `MobileSAM` and `SAM` using the same API. For more usage information, please see the [SAM page](sam.md).
|
||||
|
||||
### Automatically Build Segmentation Datasets Leveraging a Detection Model
|
||||
|
||||
To automatically annotate your dataset using the Ultralytics framework, utilize the `auto_annotate` function as demonstrated below:
|
||||
|
||||
!!! example
|
||||
|
||||
=== "Python"
|
||||
|
||||
```python
|
||||
from ultralytics.data.annotator import auto_annotate
|
||||
|
||||
auto_annotate(data="path/to/images", det_model="yolo11x.pt", sam_model="mobile_sam.pt")
|
||||
```
|
||||
|
||||
{% include "macros/sam-auto-annotate.md" %}
|
||||
|
||||
## Citations and Acknowledgements
|
||||
|
||||
If you find MobileSAM useful in your research or development work, please consider citing our paper:
|
||||
|
||||
@ -292,18 +292,7 @@ To auto-annotate your dataset using SAM 2, follow this example:
|
||||
auto_annotate(data="path/to/images", det_model="yolo11x.pt", sam_model="sam2_b.pt")
|
||||
```
|
||||
|
||||
| Argument | Type | Description | Default |
|
||||
| ------------ | ----------------------- | ------------------------------------------------------------------------------------------------------- | -------------- |
|
||||
| `data` | `str` | Path to a folder containing images to be annotated. | |
|
||||
| `det_model` | `str`, optional | Pre-trained YOLO detection model. Defaults to 'yolo11x.pt'. | `'yolo11x.pt'` |
|
||||
| `sam_model` | `str`, optional | Pre-trained SAM 2 segmentation model. Defaults to 'sam2_b.pt'. | `'sam2_b.pt'` |
|
||||
| `device` | `str`, optional | Device to run the models on. Defaults to an empty string (CPU or GPU, if available). | |
|
||||
| `conf` | `float`, optional | Confidence threshold for detection model; default is 0.25. | `0.25` |
|
||||
| `iou` | `float`, optional | IoU threshold for filtering overlapping boxes in detection results; default is 0.45. | `0.45` |
|
||||
| `imgsz` | `int`, optional | Input image resize dimension; default is 640. | `640` |
|
||||
| `max_det` | `int`, optional | Limits detections per image to control outputs in dense scenes. | `300` |
|
||||
| `classes` | `list`, optional | Filters predictions to specified class IDs, returning only relevant detections. | `None` |
|
||||
| `output_dir` | `str`, `None`, optional | Directory to save the annotated results. Defaults to a 'labels' folder in the same directory as 'data'. | `None` |
|
||||
{% include "macros/sam-auto-annotate.md" %}
|
||||
|
||||
This function facilitates the rapid creation of high-quality segmentation datasets, ideal for researchers and developers aiming to accelerate their projects.
|
||||
|
||||
|
||||
@ -208,18 +208,7 @@ To auto-annotate your dataset with the Ultralytics framework, use the `auto_anno
|
||||
auto_annotate(data="path/to/images", det_model="yolo11x.pt", sam_model="sam_b.pt")
|
||||
```
|
||||
|
||||
| Argument | Type | Description | Default |
|
||||
| ------------ | --------------------- | ------------------------------------------------------------------------------------------------------- | -------------- |
|
||||
| `data` | `str` | Path to a folder containing images to be annotated. | |
|
||||
| `det_model` | `str`, optional | Pre-trained YOLO detection model. Defaults to 'yolo11x.pt'. | `'yolo11x.pt'` |
|
||||
| `sam_model` | `str`, optional | Pre-trained SAM segmentation model. Defaults to 'sam_b.pt'. | `'sam_b.pt'` |
|
||||
| `device` | `str`, optional | Device to run the models on. Defaults to an empty string (CPU or GPU, if available). | |
|
||||
| `conf` | `float`, optional | Confidence threshold for detection model; default is 0.25. | `0.25` |
|
||||
| `iou` | `float`, optional | IoU threshold for filtering overlapping boxes in detection results; default is 0.45. | `0.45` |
|
||||
| `imgsz` | `int`, optional | Input image resize dimension; default is 640. | `640` |
|
||||
| `max_det` | `int`, optional | Limits detections per image to control outputs in dense scenes. | `300` |
|
||||
| `classes` | `list`, optional | Filters predictions to specified class IDs, returning only relevant detections. | `None` |
|
||||
| `output_dir` | `str`, None, optional | Directory to save the annotated results. Defaults to a 'labels' folder in the same directory as 'data'. | `None` |
|
||||
{% include "macros/sam-auto-annotate.md" %}
|
||||
|
||||
The `auto_annotate` function takes the path to your images, with optional arguments for specifying the pre-trained detection and SAM segmentation models, the device to run the models on, and the output directory for saving the annotated results.
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user