์ฝ˜ํ…์ธ ๋กœ ๊ฑด๋„ˆ๋›ฐ๊ธฐ

์ฐธ์กฐ ultralytics/data/annotator.py

์ฐธ๊ณ 

์ด ํŒŒ์ผ์€ https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/data/annotator .py์—์„œ ํ™•์ธํ•  ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค. ๋ฌธ์ œ๋ฅผ ๋ฐœ๊ฒฌํ•˜๋ฉด ํ’€ ๋ฆฌํ€˜์ŠคํŠธ (๐Ÿ› ๏ธ)๋ฅผ ์ œ์ถœํ•˜์—ฌ ๋ฌธ์ œ๋ฅผ ํ•ด๊ฒฐํ•˜๋„๋ก ๋„์™€์ฃผ์„ธ์š”. ๊ฐ์‚ฌํ•ฉ๋‹ˆ๋‹ค ๐Ÿ™!



ultralytics.data.annotator.auto_annotate(data, det_model='yolov8x.pt', sam_model='sam_b.pt', device='', output_dir=None)

YOLO ๊ฐ์ฒด ๊ฐ์ง€ ๋ชจ๋ธ๊ณผ SAM ์„ธ๋ถ„ํ™” ๋ชจ๋ธ์„ ์‚ฌ์šฉํ•˜์—ฌ ์ด๋ฏธ์ง€์— ์ž๋™์œผ๋กœ ์ฃผ์„์„ ๋‹ฌ ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค.

๋งค๊ฐœ๋ณ€์ˆ˜:

์ด๋ฆ„ ์œ ํ˜• ์„ค๋ช… ๊ธฐ๋ณธ๊ฐ’
data str

์ฃผ์„์„ ๋‹ฌ ์ด๋ฏธ์ง€๊ฐ€ ํฌํ•จ๋œ ํด๋”์˜ ๊ฒฝ๋กœ์ž…๋‹ˆ๋‹ค.

ํ•„์ˆ˜
det_model str

์‚ฌ์ „ ํ•™์Šต๋œ YOLO ํƒ์ง€ ๋ชจ๋ธ. ๊ธฐ๋ณธ๊ฐ’์€ 'yolov8x.pt'์ž…๋‹ˆ๋‹ค.

'yolov8x.pt'
sam_model str

์‚ฌ์ „ ํ•™์Šต๋œ SAM ์„ธ๋ถ„ํ™” ๋ชจ๋ธ. ๊ธฐ๋ณธ๊ฐ’์€ 'sam_b.pt'์ž…๋‹ˆ๋‹ค.

'sam_b.pt'
device str

๋ชจ๋ธ์„ ์‹คํ–‰ํ•  ์žฅ์น˜์ž…๋‹ˆ๋‹ค. ๊ธฐ๋ณธ๊ฐ’์€ ๋นˆ ๋ฌธ์ž์—ด(์‚ฌ์šฉ ๊ฐ€๋Šฅํ•œ ๊ฒฝ์šฐ CPU ๋˜๋Š” GPU)์ž…๋‹ˆ๋‹ค.

''
output_dir str | None | optional

๋””๋ ‰ํ„ฐ๋ฆฌ์— ์ฃผ์„์ด ๋‹ฌ๋ฆฐ ๊ฒฐ๊ณผ๋ฅผ ์ €์žฅํ•ฉ๋‹ˆ๋‹ค. ๊ธฐ๋ณธ๊ฐ’์€ '๋ฐ์ดํ„ฐ'์™€ ๊ฐ™์€ ๋””๋ ‰ํ„ฐ๋ฆฌ์— ์žˆ๋Š” '๋ ˆ์ด๋ธ”' ํด๋”์ž…๋‹ˆ๋‹ค.

None
์˜ˆ์ œ
from ultralytics.data.annotator import auto_annotate

auto_annotate(data='ultralytics/assets', det_model='yolov8n.pt', sam_model='mobile_sam.pt')
์˜ ์†Œ์Šค ์ฝ”๋“œ 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.

    Args:
        data (str): Path to a folder containing images to be annotated.
        det_model (str, optional): Pre-trained YOLO detection model. Defaults to 'yolov8x.pt'.
        sam_model (str, optional): Pre-trained SAM segmentation model. Defaults to 'sam_b.pt'.
        device (str, optional): Device to run the models on. Defaults to an empty string (CPU or GPU, if available).
        output_dir (str | None | optional): Directory to save the annotated results.
            Defaults to a 'labels' folder in the same directory as 'data'.

    Example:
        ```python
        from ultralytics.data.annotator import auto_annotate

        auto_annotate(data='ultralytics/assets', det_model='yolov8n.pt', sam_model='mobile_sam.pt')
        ```
    """
    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")





์ƒ์„ฑ๋จ 2023-11-12, ์—…๋ฐ์ดํŠธ๋จ 2023-11-25
์ž‘์„ฑ์ž: glenn-jocher (3)