当前位置：首页 > news >正文

YOLOv11-ultralytics-8.3.67部分代码阅读笔记-model.py

news 2025/9/13 13:39:23

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/licensefrom pathlib import Pathfrom 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 装饰器实现）。@propertydef 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 defaultCOCO 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 的只读属性，用于映射任务类型及其相关组件。@propertydef 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 类能够灵活适应不同的检测任务需求，同时保持与父类的兼容性。

查看全文

http://www.lryc.cn/news/544104.html