YOLOv8-ultralytics-8.2.103部分代码阅读笔记-predict.py
predict.py
ultralytics\models\yolo\detect\predict.py
目录
predict.py
1.所需的库和模块
2.class DetectionPredictor(BasePredictor):
1.所需的库和模块
# Ultralytics YOLO 🚀, AGPL-3.0 licensefrom ultralytics.engine.predictor import BasePredictor
from ultralytics.engine.results import Results
from ultralytics.utils import ops2.class DetectionPredictor(BasePredictor):
# 这段代码定义了一个名为 DetectionPredictor 的类，它继承自 BasePredictor 。这个类用于目标检测任务中的预测后处理，并返回一系列结果对象。
class DetectionPredictor(BasePredictor):# 扩展 BasePredictor 类的类，用于基于检测模型进行预测。"""A class extending the BasePredictor class for prediction based on a detection model.Example:```pythonfrom ultralytics.utils import ASSETSfrom ultralytics.models.yolo.detect import DetectionPredictorargs = dict(model="yolov8n.pt", source=ASSETS)predictor = DetectionPredictor(overrides=args)predictor.predict_cli()```"""# 定义了一个名为 postprocess 的方法，它接受三个参数。# 1.preds ：模型的原始预测输出。# 2.img ：模型输入的图像张量。# 3.orig_imgs ：原始图像的张量或列表。def postprocess(self, preds, img, orig_imgs):# 后期处理预测并返回结果对象列表。"""Post-processes predictions and returns a list of Results objects."""# 调用 ops.non_max_suppression 函数对预测结果进行非极大值抑制（NMS），以去除重叠的检测框。这个函数接受以下参数。preds = ops.non_max_suppression(# 模型的预测输出。preds,# 预测的置信度阈值。self.args.conf,# 交并比（IoU）阈值。self.args.iou,# 是否执行类别无关的NMS。agnostic=self.args.agnostic_nms,# 每个图像允许的最大检测框数量。max_det=self.args.max_det,# 要检测的类别列表。classes=self.args.classes,)# 检查 orig_imgs 是否不是列表类型。if not isinstance(orig_imgs, list): # input images are a torch.Tensor, not a list# 如果 orig_imgs 不是列表，将其从PyTorch张量转换为NumPy数组列表。# def convert_torch2numpy_batch(batch: torch.Tensor) -> np.ndarray:# -> 将 PyTorch 张量批量转换为 NumPy 数组，通常用于将模型输出的张量（例如图像数据）从 PyTorch 格式转换为 NumPy 格式，以便进行进一步处理或可视化。返回转换后的 NumPy 数组。# -> return (batch.permute(0, 2, 3, 1).contiguous() * 255).clamp(0, 255).to(torch.uint8).cpu().numpy()orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)# 初始化一个空列表 results ，用于存储结果对象。results = []# 遍历 处理后的预测结果 、 原始图像 和 图像路径 。for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):# 将预测的边界框从模型输入图像的空间尺度转换回原始图像的空间尺度。# def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None, padding=True, xywh=False):# -> 用于将边界框按照一定的比例缩放并调整到原始图像的形状。使用 clip_boxes 函数将缩放后的边界框剪辑到 原始图像 的形状内，并返回结果。# -> return clip_boxes(boxes, img0_shape)pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)# 创建一个 Results 对象，包含 原始图像 、 图像路径 、 类别名称 和 预测的边界框 ，然后添加到 results 列表中。# class Results(SimpleClass):# -> Results 类的目的是封装检测或分割模型的输出结果，包括原始图像、检测框、掩码、概率、关键点、方向边界框（Oriented Bounding Box，简称 OBB）和速度信息。# -> def __init__(self, orig_img, path, names, boxes=None, masks=None, probs=None, keypoints=None, obb=None, speed=None) -> None:results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))# 返回包含结果对象的列表。return results# 这个方法的主要功能是对模型的预测结果进行后处理，包括应用NMS、尺度转换和创建结果对象。这些结果对象可以用于进一步的分析、可视化或保存预测结果。