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:
Muhammad Rizwan Munawar 2024-12-25 17:30:12 +05:00 committed by GitHub
parent d88e923595
commit 93862d3640
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 30 additions and 24 deletions

View 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). |

View File

@ -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:

View File

@ -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.

View File

@ -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.