auto_annotate(
data,
det_model="yolov8x.pt",
sam_model="sam_b.pt",
device="",
output_dir=None,
)
Automatically annotates images using a YOLO object detection model and a SAM segmentation model.
This function processes images in a specified directory, detects objects using a YOLO model, and then generates
segmentation masks using a SAM model. The resulting annotations are saved as text files.
Parameters:
Name |
Type |
Description |
Default |
data |
str
|
Path to a folder containing images to be annotated.
|
required
|
det_model |
str
|
Path or name of the pre-trained YOLO detection model.
|
'yolov8x.pt'
|
sam_model |
str
|
Path or name of the pre-trained SAM segmentation model.
|
'sam_b.pt'
|
device |
str
|
Device to run the models on (e.g., 'cpu', 'cuda', '0').
|
''
|
output_dir |
str | None
|
Directory to save the annotated results. If None, a default directory is created.
|
None
|
Examples:
>>> from ultralytics.data.annotator import auto_annotate
>>> auto_annotate(data="ultralytics/assets", det_model="yolov8n.pt", sam_model="mobile_sam.pt")
Notes
- The function creates a new directory for output if not specified.
- Annotation results are saved as text files with the same names as the input images.
- Each line in the output text file represents a detected object with its class ID and segmentation points.
Source code in ultralytics/data/annotator.py
| def auto_annotate(data, det_model="yolov8x.pt", sam_model="sam_b.pt", device="", output_dir=None):
"""
Automatically annotates images using a YOLO object detection model and a SAM segmentation model.
This function processes images in a specified directory, detects objects using a YOLO model, and then generates
segmentation masks using a SAM model. The resulting annotations are saved as text files.
Args:
data (str): Path to a folder containing images to be annotated.
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').
output_dir (str | None): Directory to save the annotated results. If None, a default directory is created.
Examples:
>>> from ultralytics.data.annotator import auto_annotate
>>> auto_annotate(data="ultralytics/assets", det_model="yolov8n.pt", sam_model="mobile_sam.pt")
Notes:
- The function creates a new directory for output if not specified.
- Annotation results are saved as text files with the same names as the input images.
- Each line in the output text file represents a detected object with its class ID and segmentation points.
"""
det_model = YOLO(det_model)
sam_model = SAM(sam_model)
data = Path(data)
if not output_dir:
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)
for result in det_results:
class_ids = result.boxes.cls.int().tolist() # noqa
if len(class_ids):
boxes = result.boxes.xyxy # Boxes object for bbox outputs
sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)
segments = sam_results[0].masks.xyn # noqa
with open(f"{Path(output_dir) / Path(result.path).stem}.txt", "w") as f:
for i in range(len(segments)):
s = segments[i]
if len(s) == 0:
continue
segment = map(str, segments[i].reshape(-1).tolist())
f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")
|