YOLOv11-ultralytics-8.3.67部分代码阅读笔记-model.py
model.py
ultralytics\models\yolo\model.py
目录
model.py
1.所需的库和模块
2.class YOLO(Model):
3.class YOLOWorld(Model):
1.所需的库和模块
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from pathlib import Path
from ultralytics.engine.model import Model
from ultralytics.models import yolo
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, OBBModel, PoseModel, SegmentationModel, WorldModel
from ultralytics.utils import ROOT, yaml_load
2.class YOLO(Model):
# 这段代码定义了一个名为 YOLO 的类，它继承自 Model 类，并通过动态加载和初始化不同的子模块来支持多种任务（如分类、检测、分割等）。
# 定义了一个名为 YOLO 的类，继承自 Model 类。 Model 是父类，包含一些通用的模型初始化和管理功能。
# class Model(nn.Module):
# -> 继承自 torch.nn.Module ，用于封装和操作 YOLO 模型的各种功能，包括模型加载、训练、预测、导出等。
# -> def __init__(self, model: Union[str, Path] = "yolo11n.pt", task: str = None, verbose: bool = False,) -> None:
class YOLO(Model):
# YOLO（You Only Look Once）物体检测模型。
"""YOLO (You Only Look Once) object detection model."""
# 定义了 YOLO 类的初始化方法 __init__ ，它接受以下参数 ：
# 1.model ：默认值为 "yolo11n.pt" ，表示模型文件的路径或名称。
# 2.task ：表示任务类型（如分类、检测等），默认为 None 。
# 3.verbose ：布尔值，表示是否输出详细信息，默认为 False 。
def __init__(self, model="yolo11n.pt", task=None, verbose=False):
# 初始化 YOLO 模型，如果模型文件名包含‘-world’则切换到 YOLOWorld。
"""Initialize YOLO model, switching to YOLOWorld if model filename contains '-world'."""
# 使用 Path 类（可能来自 pathlib 模块）将 model 参数转换为路径对象，方便后续对文件名和扩展名的操作。
path = Path(model)
# 检查模型文件名是否包含 -world ，并且文件扩展名是否为 .pt 、 .yaml 或 .yml 。 如果满足条件，说明这是一个 YOLOWorld 模型（可能是某种特定的 YOLO 变体）。
if "-world" in path.stem and path.suffix in {".pt", ".yaml", ".yml"}: # if YOLOWorld PyTorch model
# 如果是 YOLOWorld 模型，则创建一个 YOLOWorld 类的实例 new_instance ，并将 路径 和 verbose 参数传递给它。
new_instance = YOLOWorld(path, verbose=verbose)
# 将 当前实例的类动态 更改为 YOLOWorld 的类类型。这使得 当前实例的行为和属性 与 YOLOWorld 实例一致。
self.__class__ = type(new_instance)
# 将 当前实例的属性字典 （ __dict__ ）替换为 new_instance 的属性字典。这一步确保 当前实例的属性 与 YOLOWorld 实例完全一致。
self.__dict__ = new_instance.__dict__
# 如果模型文件名不满足 YOLOWorld 的条件，则执行默认的 YOLO 初始化逻辑。
else:
# Continue with default YOLO initialization
# 调用父类 Model 的初始化方法，将 model 、 task 和 verbose 参数传递给父类。
super().__init__(model=model, task=task, verbose=verbose)
# 定义了一个名为 task_map 的属性，它是一个只读属性（通过 @property 装饰器实现）。
@property
def task_map(self):
# 将头部映射到模型、训练器、验证器和预测器类别。
"""Map head to model, trainer, validator, and predictor classes."""
# 返回一个字典，字典的键 表示不同的任务类型 （如分类、检测等），值是一个 包含模型、训练器、验证器和预测器的子字典 。
return {
# 定义了 分类任务的映射 ，包括 ：
"classify": {
# 分类模型类。
"model": ClassificationModel,
# 分类任务的训练器类。
"trainer": yolo.classify.ClassificationTrainer,
# 分类任务的验证器类。
"validator": yolo.classify.ClassificationValidator,
# 分类任务的预测器类。
"predictor": yolo.classify.ClassificationPredictor,
},
# 定义了 检测任务 的映射，包括 检测模型 、 训练器 、 验证器 和 预测器 。
"detect": {
"model": DetectionModel,
"trainer": yolo.detect.DetectionTrainer,
"validator": yolo.detect.DetectionValidator,
"predictor": yolo.detect.DetectionPredictor,
},
# 定义了 分割任务 的映射，包括分割模型、训练器、验证器和预测器。
"segment": {
"model": SegmentationModel,
"trainer": yolo.segment.SegmentationTrainer,
"validator": yolo.segment.SegmentationValidator,
"predictor": yolo.segment.SegmentationPredictor,
},
# 定义了 姿态估计 任务的映射，包括姿态模型、训练器、验证器和预测器。
"pose": {
"model": PoseModel,
"trainer": yolo.pose.PoseTrainer,
"validator": yolo.pose.PoseValidator,
"predictor": yolo.pose.PosePredictor,
},
# 定义了 定向边界框 （Oriented Bounding Box）任务的映射，包括 OBB 模型、训练器、验证器和预测器。
"obb": {
"model": OBBModel,
"trainer": yolo.obb.OBBTrainer,
"validator": yolo.obb.OBBValidator,
"predictor": yolo.obb.OBBPredictor,
},
# 结束 task_map 属性的定义。
}
# 这段代码实现了一个灵活的 YOLO 类，支持多种任务（如分类、检测、分割等）。它通过动态加载不同的子模块（如 YOLOWorld 或特定任务的模型、训练器、验证器和预测器）来实现功能扩展。这种设计使得 YOLO 类能够根据输入模型文件的类型和任务需求，动态调整其行为和属性，从而支持多种应用场景。3.class YOLOWorld(Model):
# 这段代码定义了一个名为 YOLOWorld 的类，它是 Model 类的子类，专门用于处理 YOLO 模型中的“世界检测”任务（是针对某种特定场景或数据集的检测任务）。
# 定义了一个名为 YOLOWorld 的类，继承自 Model 类。 Model 是父类，包含通用的模型初始化、训练和预测功能。
class YOLOWorld(Model):
# YOLO-World 物体检测模型。
"""YOLO-World object detection model."""
# 定义了 YOLOWorld 类的初始化方法 __init__ ，它接受以下参数 ：
# 1.model ：默认值为 "yolov8s-world.pt" ，表示模型文件的路径或名称。
# 2.verbose ：布尔值，表示是否输出详细信息，默认为 False 。
# 方法的返回类型为 None 。
def __init__(self, model="yolov8s-world.pt", verbose=False) -> None:
# 使用预训练模型文件初始化 YOLOv8-World 模型。
# 加载 YOLOv8-World 模型进行对象检测。如果未提供自定义类名，则分配默认 COCO 类名。
"""
Initialize YOLOv8-World model with a pre-trained model file.
Loads a YOLOv8-World model for object detection. If no custom class names are provided, it assigns default
COCO class names.
Args:
model (str | Path): Path to the pre-trained model file. Supports *.pt and *.yaml formats.
verbose (bool): If True, prints additional information during initialization.
"""
# 调用父类 Model 的初始化方法，将 model 、 task 和 verbose 参数传递给父类。 这里明确指定了任务类型为 "detect" ，表示这是一个用于 目标检测 的模型。
super().__init__(model=model, task="detect", verbose=verbose)
# Assign default COCO class names when there are no custom names
# 检查当前模型对象是否没有 names 属性。 names 属性通常用于 存储模型的类别名称 。
if not hasattr(self.model, "names"):
# 如果模型没有 names 属性，则从 ROOT / "cfg/datasets/coco8.yaml" 文件中加载类别名称。
# yaml_load 是一个函数（可能在其他地方定义），用于加载 YAML 文件的内容。
# .get("names") 从 YAML 文件中提取 类别名称列表 ，并将其赋值给模型的 names 属性。
# 这一步确保模型始终有默认的类别名称（如 COCO 数据集的类别）。
self.model.names = yaml_load(ROOT / "cfg/datasets/coco8.yaml").get("names")
# 定义了一个名为 task_map 的只读属性，用于映射任务类型及其相关组件。
@property
def task_map(self):
# 将头部映射到模型、验证器和预测器类别。
"""Map head to model, validator, and predictor classes."""
# 返回一个字典，字典的键是 任务类型 ，值是一个 包含模型、验证器、预测器和训练器的子字典 。
return {
# 定义了检测任务的映射。
"detect": {
# 指定检测任务的 模型类 为 WorldModel 。
"model": WorldModel,
# 指定检测任务的 验证器类 为 yolo.detect.DetectionValidator 。
"validator": yolo.detect.DetectionValidator,
# 指定检测任务的 预测器类 为 yolo.detect.DetectionPredictor 。
"predictor": yolo.detect.DetectionPredictor,
# 指定检测任务的 训练器类 为 yolo.world.WorldTrainer 。
"trainer": yolo.world.WorldTrainer,
# 结束检测任务的映射定义。
}
# 结束 task_map 属性的定义。
}
# 定义了一个名为 set_classes 的方法，用于设置模型的类别名称。 参数 ：
# 1.classes ：是一个类别名称列表。
def set_classes(self, classes):
# 设置类别。
"""
Set classes.
Args:
classes (List(str)): A list of categories i.e. ["person"].
"""
# 调用模型的 set_classes 方法，将传入的类别名称列表设置到模型中。
self.model.set_classes(classes)
# Remove background if it's given
# 定义一个变量 background ，其值为一个空格字符串。这可能是用于标识“背景”类别的占位符。
background = " "
# 检查类别列表中是否包含背景类别（空格字符串）。
if background in classes:
# 如果存在背景类别，则从类别列表中移除它。
classes.remove(background)
# 将 更新后的 类别列表 赋值给模型的 names 属性。
self.model.names = classes
# Reset method class names
# self.predictor = None # reset predictor otherwise old names remain
# 检查当前实例是否已经初始化了预测器。
if self.predictor:
# 如果预测器已初始化，则更新预测器模型的 names 属性，以确保类别名称与模型保持一致。
self.predictor.model.names = classes
# 这段代码实现了一个 YOLOWorld 类，专门用于处理目标检测任务，特别是针对某种特定场景或数据集的检测任务。它继承自通用的 Model 类，并通过以下方式扩展了功能。动态加载默认类别名称：如果模型没有定义类别名称，则从预设的 YAML 文件中加载 COCO 数据集的类别名称。任务映射：通过 task_map 属性，为检测任务定义了专门的模型、验证器、预测器和训练器。类别名称管理：通过 set_classes 方法，允许用户自定义类别名称，并确保模型和预测器的类别名称保持一致。背景类别处理：在设置类别名称时，自动移除背景类别（如果存在）。这种设计使得 YOLOWorld 类能够灵活适应不同的检测任务需求，同时保持与父类的兼容性。