YOLOv11-ultralytics-8.3.67部分代码阅读笔记-annotator.py
annotator.py
ultralytics\data\annotator.py
目录
annotator.py
1.所需的库和模块
2.def auto_annotate(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,):
1.所需的库和模块
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from pathlib import Path
from ultralytics import SAM, YOLO2.def auto_annotate(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,):
# 这段代码定义了一个名为 auto_annotate 的函数，用于自动标注数据集中的目标检测和分割结果。它结合了目标检测模型（如 YOLO）和分割模型（如 SAM），将检测到的目标转换为分割掩码，并将结果保存为标注文件。
# 定义了一个名为 auto_annotate 的函数，接受以下参数 ：
# 1.data ：输入数据的路径（可以是图像或视频文件）。
# 2.det_model ：目标检测模型的路径，默认为 "yolo11x.pt" 。
# 3.sam_model ：分割模型的路径，默认为 "sam_b.pt" 。
# 4.device ：运行模型的设备（如 "cpu" 或 "cuda" ），默认为空字符串。
# 5.conf ：检测的置信度阈值，默认为 0.25 。
# 6.iou ：非最大抑制（NMS）的 IoU 阈值，默认为 0.45 。
# 7.imgsz ：输入图像的大小，默认为 640 。
# 8.max_det ：每张图像的最大检测数，默认为 300 。
# 9.classes ：指定检测的类别，默认为 None （检测所有类别）。
# 10.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,
max_det=300,
classes=None,
output_dir=None,
):
# 使用 YOLO 物体检测模型和 SAM 分割模型自动注释图像。
# 此函数处理指定目录中的图像，使用 YOLO 模型检测物体，然后使用 SAM 模型生成分割掩码。生成的注释保存为文本文件。
# 示例：
# >>> from ultralytics.data.annotator import auto_annotate
# >>> auto_annotate(data="ultralytics/assets", det_model="yolo11n.pt", sam_model="mobile_sam.pt")
# 注意事项：
# - 如果未指定，该函数将创建一个新的输出目录。
# - 注释结果保存为与输入图像同名的文本文件。
# - 输出文本文件中的每一行代表一个检测到的对象及其类别 ID 和分割点。
"""
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').
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:
>>> from ultralytics.data.annotator import auto_annotate
>>> auto_annotate(data="ultralytics/assets", det_model="yolo11n.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.
"""
# 加载目标检测模型（如 YOLO），并将其实例化为 det_model 。
det_model = YOLO(det_model)
# 加载分割模型（如 SAM），并将其实例化为 sam_model 。
sam_model = SAM(sam_model)
# 这段代码的作用是准备自动标注的输入数据和输出目录，并调用目标检测模型对数据进行处理。
# 将输入数据路径 data 转换为 Path 对象，便于后续操作。 Path 是 pathlib 模块中的一个类，提供了方便的路径操作方法。
data = Path(data)
# 检查 是否提供了输出目录 output_dir 。如果没有提供，则根据输入数据路径动态生成一个默认的输出目录。
if not output_dir:
# 如果未指定 output_dir ，则生成默认的输出目录路径。
# data.parent ：获取输入数据的父目录。
# data.stem ：获取输入数据的文件名（不包含扩展名）。
# 默认输出目录为 ：<输入数据的父目录>/<输入数据的文件名>_auto_annotate_labels 。
output_dir = data.parent / f"{data.stem}_auto_annotate_labels"
# 使用 Path.mkdir 方法创建输出目录。
# parents=True ：如果父目录不存在，则一并创建。
# exist_ok=True ：如果目录已存在，则不会抛出异常。这一步确保输出目录存在，避免后续写入文件时出现错误。
Path(output_dir).mkdir(exist_ok=True, parents=True)
# 调用目标检测模型 det_model 对输入数据进行处理。
# data ：输入数据路径。
# stream=True ：以流式处理的方式返回检测结果，便于逐帧处理（例如处理视频或批量图像）。
# device=device ：指定运行设备（如 "cpu" 或 "cuda" ）。
# conf=conf ：设置检测的置信度阈值（默认为 0.25 ）。
# iou=iou ：设置非最大抑制（NMS）的 IoU 阈值（默认为 0.45 ）。
# imgsz=imgsz ：设置输入图像的大小（默认为 640 ）。
# max_det=max_det ：设置每张图像的最大检测数（默认为 300 ）。
# classes=classes ：指定检测的类别（默认为 None ，即检测所有类别）。
# det_results 是一个生成器，包含每张图像的检测结果。每次迭代返回一个 result 对象，包含 检测到的目标信息 （如边界框、类别 ID 等）。
det_results = det_model(
data, stream=True, device=device, conf=conf, iou=iou, imgsz=imgsz, max_det=max_det, classes=classes
)
# 这段代码的作用是。准备输入数据和输出目录：将输入数据路径转换为 Path 对象。如果未指定输出目录，则根据输入数据路径生成默认的输出目录。创建输出目录，确保其存在。调用目标检测模型：使用目标检测模型对输入数据进行处理，获取检测结果。检测结果以流式方式返回，便于逐帧处理（例如处理视频或批量图像）。这种设计使得自动标注过程更加灵活和高效，能够处理各种输入数据（如单张图像、批量图像或视频），并确保输出目录存在，便于后续保存标注结果。
# 这段代码的作用是处理目标检测模型的输出结果，并使用分割模型（如 SAM）生成分割掩码，最终将分割结果保存为标注文件。
# 遍历目标检测模型 det_model 返回的检测结果。 det_results 是一个生成器，每次迭代返回一个 result 对象，包含单张图像的检测结果。
for result in det_results:
# 提取检测到的目标类别 ID 。
# result.boxes.cls ：获取检测结果中的类别 ID（通常是一个张量）。
# .int() ：将类别 ID 转换为整数类型。
# .tolist() ：将张量转换为 Python 列表。
# class_ids 是一个列表，包含 每个检测到的目标的类别 ID 。
class_ids = result.boxes.cls.int().tolist() # noqa
# 检查 是否检测到目标 。如果 class_ids 的长度为 0，则跳过后续处理。
if len(class_ids):
# 提取 检测到的目标边界框 。
# result.boxes.xyxy ：获取边界框的坐标，格式为 (x1, y1, x2, y2) ，表示左上角和右下角的坐标。
boxes = result.boxes.xyxy # Boxes object for bbox outputs
# 使用分割模型 sam_model 对原始图像进行分割。
# result.orig_img ：原始图像。
# bboxes=boxes ：检测到的边界框，用于指导分割模型生成分割掩码。
# verbose=False ：不输出日志信息。
# save=False ：不保存分割结果到磁盘。
# device=device ：指定运行设备（如 "cpu" 或 "cuda" ）。
# sam_results 是分割模型的输出结果，包含 分割掩码等信息 。
sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)
# 提取分割结果中的 归一化掩码 。
# sam_results[0].masks.xyn ：获取归一化掩码（坐标范围为 [0, 1] ）。
# segments 是一个列表，包含每个检测到的目标的分割掩码。
segments = sam_results[0].masks.xyn # noqa
# 打开一个文件用于写入标注结果。
# 文件路径为输出目录下，以输入图像的文件名（无扩展名）加上 .txt 作为文件名。
# 使用 Path(result.path).stem 获取输入图像的文件名（无扩展名）。
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
# 将分割掩码转换为字符串列表。
# segments[i].reshape(-1) ：将掩码展平为一维数组。
# .tolist() ：将数组转换为列表。
# map(str, ...) ：将列表中的每个元素转换为字符串。
segment = map(str, segments[i].reshape(-1).tolist())
# 将类别 ID 和分割掩码写入文件，格式为 ： <class_id> <x1> <y1> <x2> <y2> ...
# class_ids[i] ：当前目标的类别 ID。
# " ".join(segment) ：将分割掩码的坐标拼接为字符串。
# \n ：换行符，用于分隔每个目标的标注信息。
f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")
# 这段代码的作用是。处理目标检测结果：提取检测到的目标类别 ID 和边界框。如果未检测到目标，则跳过后续处理。生成分割掩码：使用分割模型（如 SAM）根据检测到的边界框生成分割掩码。保存标注结果：将每个目标的类别 ID 和分割掩码保存为标注文件（如 .txt 格式）。文件路径为输出目录下，以输入图像的文件名（无扩展名）加上 .txt 作为文件名。这种设计使得自动标注过程更加高效和灵活，能够处理各种输入数据（如单张图像或视频），并生成可用于后续训练或分析的标注文件。
# 这段代码的作用是。加载模型：加载目标检测模型（如 YOLO）和分割模型（如 SAM）。处理输入数据：对输入数据进行目标检测，获取检测结果。生成分割掩码：使用分割模型将检测到的目标转换为分割掩码。保存标注结果：将分割掩码和类别 ID 保存为标注文件（如 .txt 格式）。这种设计使得自动标注过程更加高效和灵活，适用于需要大量标注数据的场景。
# def auto_annotate(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,):