Optimize Auto-Annotation with all args (#17315)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
parent
2a1fabcf83
commit
bf1d076e20
3 changed files with 19 additions and 2 deletions
|
|
@ -262,6 +262,8 @@ To auto-annotate your dataset using SAM 2, follow this example:
|
||||||
| `conf` | `float`, optional | Confidence threshold for detection model; default is 0.25. | `0.25` |
|
| `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` |
|
| `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` |
|
| `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` |
|
| `output_dir` | `str`, `None`, optional | Directory to save the annotated results. Defaults to a 'labels' folder in the same directory as 'data'. | `None` |
|
||||||
|
|
||||||
This function facilitates the rapid creation of high-quality segmentation datasets, ideal for researchers and developers aiming to accelerate their projects.
|
This function facilitates the rapid creation of high-quality segmentation datasets, ideal for researchers and developers aiming to accelerate their projects.
|
||||||
|
|
|
||||||
|
|
@ -217,6 +217,8 @@ To auto-annotate your dataset with the Ultralytics framework, use the `auto_anno
|
||||||
| `conf` | `float`, optional | Confidence threshold for detection model; default is 0.25. | `0.25` |
|
| `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` |
|
| `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` |
|
| `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` |
|
| `output_dir` | `str`, None, optional | Directory to save the annotated results. Defaults to a 'labels' folder in the same directory as 'data'. | `None` |
|
||||||
|
|
||||||
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.
|
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.
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,16 @@ from ultralytics import SAM, YOLO
|
||||||
|
|
||||||
|
|
||||||
def auto_annotate(
|
def auto_annotate(
|
||||||
data, det_model="yolo11x.pt", sam_model="sam_b.pt", device="", conf=0.25, iou=0.45, imgsz=640, output_dir=None
|
data,
|
||||||
|
det_model="yolo11x.pt",
|
||||||
|
sam_model="sam_b.pt",
|
||||||
|
device="",
|
||||||
|
conf=0.25,
|
||||||
|
iou=0.45,
|
||||||
|
imgsz=640,
|
||||||
|
max_det=300,
|
||||||
|
classes=None,
|
||||||
|
output_dir=None,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Automatically annotates images using a YOLO object detection model and a SAM segmentation model.
|
Automatically annotates images using a YOLO object detection model and a SAM segmentation model.
|
||||||
|
|
@ -22,6 +31,8 @@ def auto_annotate(
|
||||||
conf (float): Confidence threshold for detection model; default is 0.25.
|
conf (float): Confidence threshold for detection model; default is 0.25.
|
||||||
iou (float): IoU threshold for filtering overlapping boxes in detection results; default is 0.45.
|
iou (float): IoU threshold for filtering overlapping boxes in detection results; default is 0.45.
|
||||||
imgsz (int): Input image resize dimension; default is 640.
|
imgsz (int): Input image resize dimension; default is 640.
|
||||||
|
max_det (int): Limits detections per image to control outputs in dense scenes.
|
||||||
|
classes (list): Filters predictions to specified class IDs, returning only relevant detections.
|
||||||
output_dir (str | None): Directory to save the annotated results. If None, a default directory is created.
|
output_dir (str | None): Directory to save the annotated results. If None, a default directory is created.
|
||||||
|
|
||||||
Examples:
|
Examples:
|
||||||
|
|
@ -41,7 +52,9 @@ def auto_annotate(
|
||||||
output_dir = data.parent / f"{data.stem}_auto_annotate_labels"
|
output_dir = data.parent / f"{data.stem}_auto_annotate_labels"
|
||||||
Path(output_dir).mkdir(exist_ok=True, parents=True)
|
Path(output_dir).mkdir(exist_ok=True, parents=True)
|
||||||
|
|
||||||
det_results = det_model(data, stream=True, device=device, conf=conf, iou=iou, imgsz=imgsz)
|
det_results = det_model(
|
||||||
|
data, stream=True, device=device, conf=conf, iou=iou, imgsz=imgsz, max_det=max_det, classes=classes
|
||||||
|
)
|
||||||
|
|
||||||
for result in det_results:
|
for result in det_results:
|
||||||
class_ids = result.boxes.cls.int().tolist() # noqa
|
class_ids = result.boxes.cls.int().tolist() # noqa
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue