From e8743f2ac9f43d83143e6575598282a8ec55cd88 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Thu, 31 Oct 2024 17:35:26 +0500 Subject: [PATCH] Auto annotation new parameters for SAM models (#17288) Co-authored-by: UltralyticsAssistant --- docs/en/models/sam-2.md | 5 ++++- docs/en/models/sam.md | 5 ++++- ultralytics/data/annotator.py | 9 +++++++-- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/docs/en/models/sam-2.md b/docs/en/models/sam-2.md index 9083899e..de5881c4 100644 --- a/docs/en/models/sam-2.md +++ b/docs/en/models/sam-2.md @@ -256,9 +256,12 @@ To auto-annotate your dataset using SAM 2, follow this example: | 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'. | `'yolov8x.pt'` | +| `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` | | `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 c38b06e3..fe4c01bd 100644 --- a/docs/en/models/sam.md +++ b/docs/en/models/sam.md @@ -211,9 +211,12 @@ To auto-annotate your dataset with the Ultralytics framework, use the `auto_anno | 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'. | `'yolov8x.pt'` | +| `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` | | `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 3880741d..64ee9af6 100644 --- a/ultralytics/data/annotator.py +++ b/ultralytics/data/annotator.py @@ -5,7 +5,9 @@ from pathlib import Path from ultralytics import SAM, YOLO -def auto_annotate(data, det_model="yolo11x.pt", sam_model="sam_b.pt", device="", output_dir=None): +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 +): """ Automatically annotates images using a YOLO object detection model and a SAM segmentation model. @@ -17,6 +19,9 @@ def auto_annotate(data, det_model="yolo11x.pt", sam_model="sam_b.pt", device="", det_model (str): Path or name of the pre-trained YOLO detection model. sam_model (str): Path or name of the pre-trained SAM segmentation model. device (str): Device to run the models on (e.g., 'cpu', 'cuda', '0'). + 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. output_dir (str | None): Directory to save the annotated results. If None, a default directory is created. Examples: @@ -36,7 +41,7 @@ def auto_annotate(data, det_model="yolo11x.pt", sam_model="sam_b.pt", device="", 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) + det_results = det_model(data, stream=True, device=device, conf=conf, iou=iou, imgsz=imgsz) for result in det_results: class_ids = result.boxes.cls.int().tolist() # noqa