From bf1d076e20f3e169ff25f60fb10c6e05f82db353 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Sun, 3 Nov 2024 04:55:55 +0500 Subject: [PATCH] Optimize Auto-Annotation with all args (#17315) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- docs/en/models/sam-2.md | 2 ++ docs/en/models/sam.md | 2 ++ ultralytics/data/annotator.py | 17 +++++++++++++++-- 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/docs/en/models/sam-2.md b/docs/en/models/sam-2.md index de5881c42e..86059422da 100644 --- a/docs/en/models/sam-2.md +++ b/docs/en/models/sam-2.md @@ -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` | | `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` | This function facilitates the rapid creation of high-quality segmentation datasets, ideal for researchers and developers aiming to accelerate their projects. diff --git a/docs/en/models/sam.md b/docs/en/models/sam.md index fe4c01bd8b..d6f49792ea 100644 --- a/docs/en/models/sam.md +++ b/docs/en/models/sam.md @@ -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` | | `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` | 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. diff --git a/ultralytics/data/annotator.py b/ultralytics/data/annotator.py index 64ee9af6c0..fc3b8d0765 100644 --- a/ultralytics/data/annotator.py +++ b/ultralytics/data/annotator.py @@ -6,7 +6,16 @@ from ultralytics import SAM, YOLO 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. @@ -22,6 +31,8 @@ def auto_annotate( 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. 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. Examples: @@ -41,7 +52,9 @@ def auto_annotate( output_dir = data.parent / f"{data.stem}_auto_annotate_labels" 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: class_ids = result.boxes.cls.int().tolist() # noqa