Yolov8 源码解析（三十一）
.\yolov8\ultralytics\engine\results.py
# Ultralytics YOLO 🚀, AGPL-3.0 license
"""
Ultralytics Results, Boxes and Masks classes for handling inference results.
Usage: See https://docs.ultralytics.com/modes/predict/
"""
from copy import deepcopy # 导入深拷贝函数deepcopy
from functools import lru_cache # 导入LRU缓存函数lru_cache
from pathlib import Path # 导入处理路径的Path模块
import numpy as np # 导入NumPy库
import torch # 导入PyTorch库
from ultralytics.data.augment import LetterBox # 导入augment模块中的LetterBox类
from ultralytics.utils import LOGGER, SimpleClass, ops # 导入utils模块中的LOGGER、SimpleClass和ops
from ultralytics.utils.plotting import Annotator, colors, save_one_box # 导入plotting模块中的Annotator、colors和save_one_box函数
from ultralytics.utils.torch_utils import smart_inference_mode # 导入torch_utils模块中的smart_inference_mode函数
class BaseTensor(SimpleClass):
"""
Base tensor class with additional methods for easy manipulation and device handling.
Attributes:
data (torch.Tensor | np.ndarray): Prediction data such as bounding boxes, masks, or keypoints.
orig_shape (Tuple[int, int]): Original shape of the image, typically in the format (height, width).
Methods:
cpu: Return a copy of the tensor stored in CPU memory.
numpy: Returns a copy of the tensor as a numpy array.
cuda: Moves the tensor to GPU memory, returning a new instance if necessary.
to: Return a copy of the tensor with the specified device and dtype.
Examples:
>>> import torch
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> orig_shape = (720, 1280)
>>> base_tensor = BaseTensor(data, orig_shape)
>>> cpu_tensor = base_tensor.cpu()
>>> numpy_array = base_tensor.numpy()
>>> gpu_tensor = base_tensor.cuda()
"""
def __init__(self, data, orig_shape) -> None:
"""
Initialize BaseTensor with prediction data and the original shape of the image.
Args:
data (torch.Tensor | np.ndarray): Prediction data such as bounding boxes, masks, or keypoints.
orig_shape (Tuple[int, int]): Original shape of the image in (height, width) format.
Examples:
>>> import torch
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> orig_shape = (720, 1280)
>>> base_tensor = BaseTensor(data, orig_shape)
"""
assert isinstance(data, (torch.Tensor, np.ndarray)), "data must be torch.Tensor or np.ndarray"
self.data = data # 设置BaseTensor类的data属性为传入的data
self.orig_shape = orig_shape # 设置BaseTensor类的orig_shape属性为传入的orig_shape
@property
def shape(self):
"""
Returns the shape of the underlying data tensor.
Returns:
(Tuple[int, ...]): The shape of the data tensor.
Examples:
>>> data = torch.rand(100, 4)
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
>>> print(base_tensor.shape)
(100, 4)
"""
return self.data.shape # 返回BaseTensor类的data属性的形状
def cpu(self):
"""
Returns a copy of the tensor stored in CPU memory.
Returns:
(BaseTensor): A new BaseTensor object with the data tensor moved to CPU memory.
Examples:
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]]).cuda()
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
>>> cpu_tensor = base_tensor.cpu()
>>> isinstance(cpu_tensor, BaseTensor)
True
>>> cpu_tensor.data.device
device(type='cpu')
"""
# 如果数据已经是 numpy 数组，则直接返回 self
return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.cpu(), self.orig_shape)
def numpy(self):
"""
Returns a copy of the tensor as a numpy array.
Returns:
(np.ndarray): A numpy array containing the same data as the original tensor.
Examples:
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> orig_shape = (720, 1280)
>>> base_tensor = BaseTensor(data, orig_shape)
>>> numpy_array = base_tensor.numpy()
>>> print(type(numpy_array))
<class 'numpy.ndarray'>
"""
# 如果数据已经是 numpy 数组，则直接返回 self
return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.numpy(), self.orig_shape)
def cuda(self):
"""
Moves the tensor to GPU memory.
Returns:
(BaseTensor): A new BaseTensor instance with the data moved to GPU memory if it's not already a
numpy array, otherwise returns self.
Examples:
>>> import torch
>>> from ultralytics.engine.results import BaseTensor
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
>>> gpu_tensor = base_tensor.cuda()
>>> print(gpu_tensor.data.device)
cuda:0
"""
# 将数据转换为 tensor，并移动到 GPU，然后创建新的 BaseTensor 实例返回
return self.__class__(torch.as_tensor(self.data).cuda(), self.orig_shape)
def to(self, *args, **kwargs):
"""
Return a copy of the tensor with the specified device and dtype.
Args:
*args (Any): Variable length argument list to be passed to torch.Tensor.to().
**kwargs (Any): Arbitrary keyword arguments to be passed to torch.Tensor.to().
Returns:
(BaseTensor): A new BaseTensor instance with the data moved to the specified device and/or dtype.
Examples:
>>> base_tensor = BaseTensor(torch.randn(3, 4), orig_shape=(480, 640))
>>> cuda_tensor = base_tensor.to('cuda')
>>> float16_tensor = base_tensor.to(dtype=torch.float16)
"""
# 将数据转换为 tensor，并按照参数指定的设备和数据类型进行转换，然后创建新的 BaseTensor 实例返回
return self.__class__(torch.as_tensor(self.data).to(*args, **kwargs), self.orig_shape)
# 重写 len(results)，返回基础数据张量的长度。
def __len__(self): # override len(results)
"""
Returns the length of the underlying data tensor.
Returns:
(int): The number of elements in the first dimension of the data tensor.
Examples:
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
>>> len(base_tensor)
2
"""
# 返回数据张量第一维的元素个数
return len(self.data)
# 获取指定索引的数据，返回一个包含指定索引数据的新的 BaseTensor 实例
def __getitem__(self, idx):
"""
Returns a new BaseTensor instance containing the specified indexed elements of the data tensor.
Args:
idx (int | List[int] | torch.Tensor): Index or indices to select from the data tensor.
Returns:
(BaseTensor): A new BaseTensor instance containing the indexed data.
Examples:
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
>>> result = base_tensor[0] # Select the first row
>>> print(result.data)
tensor([1, 2, 3])
"""
# 返回包含指定索引数据的新 BaseTensor 实例
return self.__class__(self.data[idx], self.orig_shape)
# Results 类用于存储和操作推理结果，继承自 SimpleClass 类
class Results(SimpleClass):
"""
A class for storing and manipulating inference results.
This class encapsulates the functionality for handling detection, segmentation, pose estimation,
and classification results from YOLO models.
Attributes:
orig_img (numpy.ndarray): Original image as a numpy array. 原始图像的 numpy 数组表示
orig_shape (Tuple[int, int]): Original image shape in (height, width) format. 原始图像的高度和宽度
boxes (Boxes | None): Object containing detection bounding boxes. 检测框的对象，可能为 None
masks (Masks | None): Object containing detection masks. 检测到的掩模的对象，可能为 None
probs (Probs | None): Object containing class probabilities for classification tasks. 分类任务的类别概率的对象，可能为 None
keypoints (Keypoints | None): Object containing detected keypoints for each object. 每个对象检测到的关键点的对象，可能为 None
obb (OBB | None): Object containing oriented bounding boxes. 方向边界框的对象，可能为 None
speed (Dict[str, float | None]): Dictionary of preprocess, inference, and postprocess speeds.
预处理、推理和后处理速度的字典，包含字符串键和浮点数或 None 值
names (Dict[int, str]): Dictionary mapping class IDs to class names. 将类别 ID 映射到类别名称的字典
path (str): Path to the image file. 图像文件的路径
_keys (Tuple[str, ...]): Tuple of attribute names for internal use. 内部使用的属性名称元组
Methods:
update: Updates object attributes with new detection results. 使用新的检测结果更新对象属性
cpu: Returns a copy of the Results object with all tensors on CPU memory. 返回所有张量在 CPU 内存上的 Results 对象副本
numpy: Returns a copy of the Results object with all tensors as numpy arrays. 返回所有张量作为 numpy 数组的 Results 对象副本
cuda: Returns a copy of the Results object with all tensors on GPU memory. 返回所有张量在 GPU 内存上的 Results 对象副本
to: Returns a copy of the Results object with tensors on a specified device and dtype. 返回指定设备和数据类型上的张量的 Results 对象副本
new: Returns a new Results object with the same image, path, and names. 返回具有相同图像、路径和名称的新 Results 对象
plot: Plots detection results on an input image, returning an annotated image. 在输入图像上绘制检测结果，返回带注释的图像
show: Shows annotated results on screen. 在屏幕上显示带注释的结果
save: Saves annotated results to file. 将带注释的结果保存到文件
verbose: Returns a log string for each task, detailing detections and classifications. 返回每个任务的日志字符串，详细描述检测和分类
save_txt: Saves detection results to a text file. 将检测结果保存到文本文件
save_crop: Saves cropped detection images. 保存裁剪后的检测图像
tojson: Converts detection results to JSON format. 将检测结果转换为 JSON 格式
Examples:
>>> results = model("path/to/image.jpg")
>>> for result in results:
... print(result.boxes) # Print detection boxes 打印检测框
... result.show() # Display the annotated image 显示带注释的图像
... result.save(filename='result.jpg') # Save annotated image 保存带注释的图像
"""
def __init__(
self, orig_img, path, names, boxes=None, masks=None, probs=None, keypoints=None, obb=None, speed=None
"""
Initialize the Results class for storing and manipulating inference results.
Args:
orig_img (numpy.ndarray): The original image as a numpy array.
path (str): The path to the image file.
names (Dict): A dictionary of class names.
boxes (torch.Tensor | None): A 2D tensor of bounding box coordinates for each detection.
masks (torch.Tensor | None): A 3D tensor of detection masks, where each mask is a binary image.
probs (torch.Tensor | None): A 1D tensor of probabilities of each class for classification task.
keypoints (torch.Tensor | None): A 2D tensor of keypoint coordinates for each detection.
obb (torch.Tensor | None): A 2D tensor of oriented bounding box coordinates for each detection.
speed (Dict | None): A dictionary containing preprocess, inference, and postprocess speeds (ms/image).
Examples:
>>> results = model("path/to/image.jpg")
>>> result = results[0] # Get the first result
>>> boxes = result.boxes # Get the boxes for the first result
>>> masks = result.masks # Get the masks for the first result
Notes:
For the default pose model, keypoint indices for human body pose estimation are:
0: Nose, 1: Left Eye, 2: Right Eye, 3: Left Ear, 4: Right Ear
5: Left Shoulder, 6: Right Shoulder, 7: Left Elbow, 8: Right Elbow
9: Left Wrist, 10: Right Wrist, 11: Left Hip, 12: Right Hip
13: Left Knee, 14: Right Knee, 15: Left Ankle, 16: Right Ankle
"""
# 存储原始图像的 numpy 数组
self.orig_img = orig_img
# 存储原始图像的形状（高度和宽度）
self.orig_shape = orig_img.shape[:2]
# 如果提供了边界框数据，则用边界框数据初始化 Boxes 类，否则为 None
self.boxes = Boxes(boxes, self.orig_shape) if boxes is not None else None # native size boxes
# 如果提供了掩码数据，则用掩码数据初始化 Masks 类，否则为 None
self.masks = Masks(masks, self.orig_shape) if masks is not None else None # native size or imgsz masks
# 如果提供了概率数据，则用概率数据初始化 Probs 类，否则为 None
self.probs = Probs(probs) if probs is not None else None
# 如果提供了关键点数据，则用关键点数据初始化 Keypoints 类，否则为 None
self.keypoints = Keypoints(keypoints, self.orig_shape) if keypoints is not None else None
# 如果提供了方向边界框数据，则用方向边界框数据初始化 OBB 类，否则为 None
self.obb = OBB(obb, self.orig_shape) if obb is not None else None
# 如果提供了速度数据，则使用提供的速度数据，否则初始化为空字典
self.speed = speed if speed is not None else {"preprocess": None, "inference": None, "postprocess": None}
# 存储类别名称的字典
self.names = names
# 存储图像文件的路径
self.path = path
# 初始化保存目录为空
self.save_dir = None
# 存储需要公开的属性名称的元组
self._keys = "boxes", "masks", "probs", "keypoints", "obb"
def __getitem__(self, idx):
"""
Return a Results object for a specific index of inference results.
Args:
idx (int | slice): Index or slice to retrieve from the Results object.
Returns:
(Results): A new Results object containing the specified subset of inference results.
Examples:
>>> results = model('path/to/image.jpg') # Perform inference
>>> single_result = results[0] # Get the first result
>>> subset_results = results[1:4] # Get a slice of results
"""
# 调用内部方法 _apply，用于处理索引操作
return self._apply("__getitem__", idx)
def __len__(self):
"""
Return the number of detections in the Results object.
Returns:
(int): The number of detections, determined by the length of the first non-empty attribute
(boxes, masks, probs, keypoints, or obb).
Examples:
>>> results = Results(orig_img, path, names, boxes=torch.rand(5, 4))
>>> len(results)
5
"""
# 遍历 self._keys 中的属性，找到第一个非空属性并返回其长度
for k in self._keys:
v = getattr(self, k)
if v is not None:
return len(v)
def update(self, boxes=None, masks=None, probs=None, obb=None):
"""
Updates the Results object with new detection data.
This method allows updating the boxes, masks, probabilities, and oriented bounding boxes (OBB) of the
Results object. It ensures that boxes are clipped to the original image shape.
Args:
boxes (torch.Tensor | None): A tensor of shape (N, 6) containing bounding box coordinates and
confidence scores. The format is (x1, y1, x2, y2, conf, class).
masks (torch.Tensor | None): A tensor of shape (N, H, W) containing segmentation masks.
probs (torch.Tensor | None): A tensor of shape (num_classes,) containing class probabilities.
obb (torch.Tensor | None): A tensor of shape (N, 5) containing oriented bounding box coordinates.
Examples:
>>> results = model('image.jpg')
>>> new_boxes = torch.tensor([[100, 100, 200, 200, 0.9, 0]])
>>> results[0].update(boxes=new_boxes)
"""
# 如果参数不为 None，则更新相应的属性值
if boxes is not None:
# 更新 boxes 属性，确保边界框被剪裁到原始图像形状内
self.boxes = Boxes(ops.clip_boxes(boxes, self.orig_shape), self.orig_shape)
if masks is not None:
# 更新 masks 属性
self.masks = Masks(masks, self.orig_shape)
if probs is not None:
# 更新 probs 属性
self.probs = probs
if obb is not None:
# 更新 obb 属性
self.obb = OBB(obb, self.orig_shape)
def _apply(self, fn, *args, **kwargs):
"""
Applies a function to all non-empty attributes and returns a new Results object with modified attributes.
This method is internally called by methods like .to(), .cuda(), .cpu(), etc.
Args:
fn (str): The name of the function to apply.
*args (Any): Variable length argument list to pass to the function.
**kwargs (Any): Arbitrary keyword arguments to pass to the function.
Returns:
(Results): A new Results object with attributes modified by the applied function.
Examples:
>>> results = model("path/to/image.jpg")
>>> for result in results:
... result_cuda = result.cuda()
... result_cpu = result.cpu()
"""
# 创建一个新的 Results 对象，用于存储应用函数后的结果
r = self.new()
# 遍历当前对象的所有属性名
for k in self._keys:
# 获取当前属性的值
v = getattr(self, k)
# 如果属性值不为 None，则对其调用指定的函数，并将结果设置到新的 Results 对象中对应的属性
if v is not None:
setattr(r, k, getattr(v, fn)(*args, **kwargs))
# 返回经过函数应用后的新 Results 对象
return r
def cpu(self):
"""
Returns a copy of the Results object with all its tensors moved to CPU memory.
This method creates a new Results object with all tensor attributes (boxes, masks, probs, keypoints, obb)
transferred to CPU memory. It's useful for moving data from GPU to CPU for further processing or saving.
Returns:
(Results): A new Results object with all tensor attributes on CPU memory.
Examples:
>>> results = model('path/to/image.jpg') # Perform inference
>>> cpu_result = results[0].cpu() # Move the first result to CPU
>>> print(cpu_result.boxes.device) # Output: cpu
"""
# 调用 _apply 方法，将所有张量移到 CPU 上
return self._apply("cpu")
def numpy(self):
"""
Converts all tensors in the Results object to numpy arrays.
Returns:
(Results): A new Results object with all tensors converted to numpy arrays.
Examples:
>>> results = model('path/to/image.jpg')
>>> numpy_result = results[0].numpy()
>>> type(numpy_result.boxes.data)
<class 'numpy.ndarray'>
Notes:
This method creates a new Results object, leaving the original unchanged. It's useful for
interoperability with numpy-based libraries or when CPU-based operations are required.
"""
# 调用 _apply 方法，将所有张量转换为 numpy 数组
return self._apply("numpy")
def cuda(self):
"""
Moves all tensors in the Results object to GPU memory.
Returns:
(Results): A new Results object with all tensors moved to CUDA device.
Examples:
>>> results = model("path/to/image.jpg")
>>> cuda_results = results[0].cuda() # Move first result to GPU
>>> for result in results:
... result_cuda = result.cuda() # Move each result to GPU
"""
# 调用 _apply 方法，将所有张量移到 GPU 上
return self._apply("cuda")
def to(self, *args, **kwargs):
"""
Moves all tensors in the Results object to the specified device and dtype.
Args:
*args (Any): Variable length argument list to be passed to torch.Tensor.to().
**kwargs (Any): Arbitrary keyword arguments to be passed to torch.Tensor.to().
Returns:
(Results): A new Results object with all tensors moved to the specified device and dtype.
Examples:
>>> results = model("path/to/image.jpg")
>>> result_cuda = results[0].to("cuda") # Move first result to GPU
>>> result_cpu = results[0].to("cpu") # Move first result to CPU
>>> result_half = results[0].to(dtype=torch.float16) # Convert first result to half precision
"""
# 调用私有方法 _apply()，将指定参数传递给它
return self._apply("to", *args, **kwargs)
def new(self):
"""
Creates a new Results object with the same image, path, names, and speed attributes.
Returns:
(Results): A new Results object with copied attributes from the original instance.
Examples:
>>> results = model("path/to/image.jpg")
>>> new_result = results[0].new()
"""
# 使用当前实例的属性创建一个新的 Results 对象并返回
return Results(orig_img=self.orig_img, path=self.path, names=self.names, speed=self.speed)
def plot(
self,
conf=True,
line_width=None,
font_size=None,
font="Arial.ttf",
pil=False,
img=None,
im_gpu=None,
kpt_radius=5,
kpt_line=True,
labels=True,
boxes=True,
masks=True,
probs=True,
show=False,
save=False,
filename=None,
):
"""
Plotting method for visualizing inference results on an image.
Args:
conf (bool): Whether to display confidence scores.
line_width (int or None): Width of lines to draw (default: None).
font_size (int or None): Size of the font for labels (default: None).
font (str): Font type for labels (default: "Arial.ttf").
pil (bool): Whether to use PIL for plotting (default: False).
img (PIL.Image.Image or None): Optional PIL image to plot results on.
im_gpu (torch.Tensor or None): Optional GPU tensor image to plot results on.
kpt_radius (int): Radius of keypoints markers (default: 5).
kpt_line (bool): Whether to draw lines between keypoints (default: True).
labels (bool): Whether to display class labels (default: True).
boxes (bool): Whether to display bounding boxes (default: True).
masks (bool): Whether to display masks (default: True).
probs (bool): Whether to display probabilities/confidence scores (default: True).
show (bool): Whether to display the plot (default: False).
save (bool): Whether to save the plot to a file (default: False).
filename (str or None): Optional filename to save the plot as.
Returns:
None
Notes:
This method allows visualization of object detection results on an image. It supports various options
such as displaying bounding boxes, masks, keypoints, labels, and confidence scores.
Examples:
>>> results = model('path/to/image.jpg')
>>> results[0].plot(show=True) # Plot and display the first result
"""
def show(self, *args, **kwargs):
"""
Display the image with annotated inference results.
This method plots the detection results on the original image and displays it. It's a convenient way to
visualize the model's predictions directly.
Args:
*args (Any): Variable length argument list to be passed to the `plot()` method.
**kwargs (Any): Arbitrary keyword arguments to be passed to the `plot()` method.
Returns:
None
Examples:
>>> results = model('path/to/image.jpg')
>>> results[0].show() # Display the first result
>>> for result in results:
... result.show() # Display all results
"""
# 调用当前对象的 plot 方法，并传递参数 show=True 及其他参数
self.plot(show=True, *args, **kwargs)
def save(self, filename=None, *args, **kwargs):
"""
Saves annotated inference results image to file.
This method plots the detection results on the original image and saves the annotated image to a file. It
utilizes the `plot` method to generate the annotated image and then saves it to the specified filename.
Args:
filename (str | Path | None): The filename to save the annotated image. If None, a default filename
is generated based on the original image path.
*args (Any): Variable length argument list to be passed to the `plot` method.
**kwargs (Any): Arbitrary keyword arguments to be passed to the `plot` method.
Examples:
>>> results = model('path/to/image.jpg')
>>> for result in results:
... result.save('annotated_image.jpg')
>>> # Or with custom plot arguments
>>> for result in results:
... result.save('annotated_image.jpg', conf=False, line_width=2)
"""
# 如果没有指定文件名，基于原始图像路径生成默认文件名
if not filename:
filename = f"results_{Path(self.path).name}"
# 调用 plot 方法生成带有标注的图像，并保存到指定的文件名
self.plot(save=True, filename=filename, *args, **kwargs)
# 返回保存的文件名
return filename
def verbose(self):
"""
Returns a log string for each task in the results, detailing detection and classification outcomes.
This method generates a human-readable string summarizing the detection and classification results. It includes
the number of detections for each class and the top probabilities for classification tasks.
Returns:
(str): A formatted string containing a summary of the results. For detection tasks, it includes the
number of detections per class. For classification tasks, it includes the top 5 class probabilities.
Examples:
>>> results = model('path/to/image.jpg')
>>> for result in results:
... print(result.verbose())
2 persons, 1 car, 3 traffic lights,
dog 0.92, cat 0.78, horse 0.64,
Notes:
- If there are no detections, the method returns "(no detections), " for detection tasks.
- For classification tasks, it returns the top 5 class probabilities and their corresponding class names.
- The returned string is comma-separated and ends with a comma and a space.
"""
# 初始化日志字符串
log_string = ""
# 获取概率和边界框信息
probs = self.probs
boxes = self.boxes
# 如果没有检测到结果，根据情况返回空字符串或者“(no detections), ”
if len(self) == 0:
return log_string if probs is None else f"{log_string}(no detections), "
# 如果有概率信息，将前五个概率最高的类别及其概率加入日志字符串
if probs is not None:
log_string += f"{', '.join(f'{self.names[j]} {probs.data[j]:.2f}' for j in probs.top5)}, "
# 如果有边界框信息，计算每个类别的检测数量，并加入日志字符串
if boxes:
for c in boxes.cls.unique():
n = (boxes.cls == c).sum() # detections per class
log_string += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "
# 返回生成的日志字符串
return log_string
# 将检测结果保存到文本文件中
def save_txt(self, txt_file, save_conf=False):
"""
Save detection results to a text file.
Args:
txt_file (str | Path): Path to the output text file.
save_conf (bool): Whether to include confidence scores in the output.
Returns:
(str): Path to the saved text file.
Examples:
>>> from ultralytics import YOLO
>>> model = YOLO('yolov8n.pt')
>>> results = model("path/to/image.jpg")
>>> for result in results:
... result.save_txt("output.txt")
Notes:
- The file will contain one line per detection or classification with the following structure:
- For detections: `class confidence x_center y_center width height`
- For classifications: `confidence class_name`
- For masks and keypoints, the specific formats will vary accordingly.
- The function will create the output directory if it does not exist.
- If save_conf is False, the confidence scores will be excluded from the output.
- Existing contents of the file will not be overwritten; new results will be appended.
"""
# 检查是否使用旋转边界框
is_obb = self.obb is not None
# 获取边界框、掩模、概率、关键点的数据
boxes = self.obb if is_obb else self.boxes
masks = self.masks
probs = self.probs
kpts = self.keypoints
# 保存文本结果的列表
texts = []
# 如果有概率信息
if probs is not None:
# 对每个类别按置信度保存
[texts.append(f"{probs.data[j]:.2f} {self.names[j]}") for j in probs.top5]
# 如果有边界框信息
elif boxes:
# 对每个边界框进行处理
for j, d in enumerate(boxes):
# 获取类别、置信度、ID
c, conf, id = int(d.cls), float(d.conf), None if d.id is None else int(d.id.item())
# 构造文本行
line = (c, *(d.xyxyxyxyn.view(-1) if is_obb else d.xywhn.view(-1)))
# 如果有掩模信息
if masks:
# 处理掩模数据
seg = masks[j].xyn[0].copy().reshape(-1) # reversed mask.xyn, (n,2) to (n*2)
line = (c, *seg)
# 如果有关键点信息
if kpts is not None:
kpt = torch.cat((kpts[j].xyn, kpts[j].conf[..., None]), 2) if kpts[j].has_visible else kpts[j].xyn
line += (*kpt.reshape(-1).tolist(),)
# 根据是否保存置信度，添加到文本行末尾
line += (conf,) * save_conf + (() if id is None else (id,))
# 格式化文本行并添加到文本列表
texts.append(("%g " * len(line)).rstrip() % line)
# 如果有文本结果
if texts:
# 确保输出目录存在
Path(txt_file).parent.mkdir(parents=True, exist_ok=True) # make directory
# 将文本行写入文件
with open(txt_file, "a") as f:
f.writelines(text + "\n" for text in texts)
# 将检测到的物体裁剪图像保存到指定目录
# 如果 self.probs 不为 None，输出警告信息并返回，因为不支持分类任务
if self.probs is not None:
LOGGER.warning("WARNING ⚠️ Classify task do not support `save_crop`.")
return
# 如果 self.obb 不为 None，输出警告信息并返回，因为不支持有向边界框任务
if self.obb is not None:
LOGGER.warning("WARNING ⚠️ OBB task do not support `save_crop`.")
return
# 遍历每个检测到的框
for d in self.boxes:
# 调用 save_one_box 函数保存单个框的裁剪图像
# 使用 d.xyxy 作为边界框坐标
# 使用 self.orig_img 的副本进行裁剪，避免修改原始图像
# 生成的文件路径为 save_dir/class_name/file_name.jpg
save_one_box(
d.xyxy,
self.orig_img.copy(),
file=Path(save_dir) / self.names[int(d.cls)] / f"{Path(file_name)}.jpg",
BGR=True,
)
def tojson(self, normalize=False, decimals=5):
"""
Converts detection results to JSON format.
This method serializes the detection results into a JSON-compatible format. It includes information
about detected objects such as bounding boxes, class names, confidence scores, and optionally
segmentation masks and keypoints.
Args:
normalize (bool): Whether to normalize the bounding box coordinates by the image dimensions.
If True, coordinates will be returned as float values between 0 and 1. Defaults to False.
decimals (int): Number of decimal places to round the output values to. Defaults to 5.
Returns:
(str): A JSON string containing the serialized detection results.
Examples:
>>> results = model("path/to/image.jpg")
>>> json_result = results[0].tojson()
>>> print(json_result)
Notes:
- For classification tasks, the JSON will contain class probabilities instead of bounding boxes.
- For object detection tasks, the JSON will include bounding box coordinates, class names, and
confidence scores.
- If available, segmentation masks and keypoints will also be included in the JSON output.
- The method uses the `summary` method internally to generate the data structure before
converting it to JSON.
"""
import json # 导入 JSON 模块
# 调用对象的 `summary` 方法生成数据结构，并转换成 JSON 格式字符串，缩进为 2
return json.dumps(self.summary(normalize=normalize, decimals=decimals), indent=2)
# 定义一个 Boxes 类，继承自 BaseTensor 类，用于管理和操作检测框。
"""
A class for managing and manipulating detection boxes.
This class provides functionality for handling detection boxes, including their coordinates, confidence scores,
class labels, and optional tracking IDs. It supports various box formats and offers methods for easy manipulation
and conversion between different coordinate systems.
Attributes:
data (torch.Tensor | numpy.ndarray): The raw tensor containing detection boxes and associated data.
orig_shape (Tuple[int, int]): The original image dimensions (height, width).
is_track (bool): Indicates whether tracking IDs are included in the box data.
xyxy (torch.Tensor | numpy.ndarray): Boxes in [x1, y1, x2, y2] format.
conf (torch.Tensor | numpy.ndarray): Confidence scores for each box.
cls (torch.Tensor | numpy.ndarray): Class labels for each box.
id (torch.Tensor | numpy.ndarray): Tracking IDs for each box (if available).
xywh (torch.Tensor | numpy.ndarray): Boxes in [x, y, width, height] format.
xyxyn (torch.Tensor | numpy.ndarray): Normalized [x1, y1, x2, y2] boxes relative to orig_shape.
xywhn (torch.Tensor | numpy.ndarray): Normalized [x, y, width, height] boxes relative to orig_shape.
Methods:
cpu(): Returns a copy of the object with all tensors on CPU memory.
numpy(): Returns a copy of the object with all tensors as numpy arrays.
cuda(): Returns a copy of the object with all tensors on GPU memory.
to(*args, **kwargs): Returns a copy of the object with tensors on specified device and dtype.
Examples:
>>> import torch
>>> boxes_data = torch.tensor([[100, 50, 150, 100, 0.9, 0], [200, 150, 300, 250, 0.8, 1]])
>>> orig_shape = (480, 640) # height, width
>>> boxes = Boxes(boxes_data, orig_shape)
>>> print(boxes.xyxy)
>>> print(boxes.conf)
>>> print(boxes.cls)
>>> print(boxes.xywhn)
"""
def __init__(self, boxes, orig_shape) -> None:
"""
Initialize the Boxes class with detection box data and the original image shape.
This class manages detection boxes, providing easy access and manipulation of box coordinates,
confidence scores, class identifiers, and optional tracking IDs. It supports multiple formats
for box coordinates, including both absolute and normalized forms.
Args:
boxes (torch.Tensor | np.ndarray): A tensor or numpy array with detection boxes of shape
(num_boxes, 6) or (num_boxes, 7). Columns should contain
[x1, y1, x2, y2, confidence, class, (optional) track_id].
orig_shape (Tuple[int, int]): The original image shape as (height, width). Used for normalization.
Attributes:
data (torch.Tensor): The raw tensor containing detection boxes and their associated data.
orig_shape (Tuple[int, int]): The original image size, used for normalization.
is_track (bool): Indicates whether tracking IDs are included in the box data.
Examples:
>>> import torch
>>> boxes = torch.tensor([[100, 50, 150, 100, 0.9, 0]])
>>> orig_shape = (480, 640)
>>> detection_boxes = Boxes(boxes, orig_shape)
>>> print(detection_boxes.xyxy)
tensor([[100., 50., 150., 100.]])
"""
# 如果输入的 boxes 是一维的，则将其转换为二维的
if boxes.ndim == 1:
boxes = boxes[None, :]
# 检查 boxes 的最后一个维度是否为 6 或 7，分别对应 xyxy, track_id, conf, cls 的不同格式
n = boxes.shape[-1]
assert n in {6, 7}, f"expected 6 or 7 values but got {n}" # xyxy, track_id, conf, cls
# 调用父类的初始化方法，将 boxes 和 orig_shape 传入
super().__init__(boxes, orig_shape)
# 设置是否包含 track_id 的标志位
self.is_track = n == 7
# 设置原始图像的形状信息
self.orig_shape = orig_shape
@property
def xyxy(self):
"""
Returns bounding boxes in [x1, y1, x2, y2] format.
Returns:
(torch.Tensor | numpy.ndarray): A tensor or numpy array of shape (n, 4) containing bounding box
coordinates in [x1, y1, x2, y2] format, where n is the number of boxes.
Examples:
>>> results = model('image.jpg')
>>> boxes = results[0].boxes
>>> xyxy = boxes.xyxy
>>> print(xyxy)
"""
# 返回数据中的前四列，即 [x1, y1, x2, y2] 形式的边界框坐标
return self.data[:, :4]
@property
def conf(self):
"""
Returns the confidence scores for each detection box.
Returns:
(torch.Tensor | numpy.ndarray): A 1D tensor or array containing confidence scores for each detection,
with shape (N,) where N is the number of detections.
Examples:
>>> boxes = Boxes(torch.tensor([[10, 20, 30, 40, 0.9, 0]]), orig_shape=(100, 100))
>>> conf_scores = boxes.conf
>>> print(conf_scores)
tensor([0.9000])
"""
# 返回数据中倒数第二列，即检测框的置信度分数
return self.data[:, -2]
@property
这段代码是一个 Python 类的初始化方法和其属性方法的定义。初始化方法用于设置检测框数据和原始图像形状，而属性方法分别用于获取边界框坐标和置信度分数。
def cls(self):
"""
Returns the class ID tensor representing category predictions for each bounding box.
Returns:
(torch.Tensor | numpy.ndarray): A tensor or numpy array containing the class IDs for each detection box.
The shape is (N,), where N is the number of boxes.
Examples:
>>> results = model('image.jpg')
>>> boxes = results[0].boxes
>>> class_ids = boxes.cls
>>> print(class_ids) # tensor([0., 2., 1.])
"""
# 返回每个检测框的类别ID张量，这些ID表示每个框所属的类别
return self.data[:, -1]
@property
def id(self):
"""
Returns the tracking IDs for each detection box if available.
Returns:
(torch.Tensor | None): A tensor containing tracking IDs for each box if tracking is enabled,
otherwise None. Shape is (N,) where N is the number of boxes.
Examples:
>>> results = model.track('path/to/video.mp4')
>>> for result in results:
... boxes = result.boxes
... if boxes.is_track:
... track_ids = boxes.id
... print(f"Tracking IDs: {track_ids}")
... else:
... print("Tracking is not enabled for these boxes.")
Notes:
- This property is only available when tracking is enabled (i.e., when `is_track` is True).
- The tracking IDs are typically used to associate detections across multiple frames in video analysis.
"""
# 如果启用了跟踪（即 `is_track` 为True），返回每个检测框的跟踪ID张量；否则返回None
return self.data[:, -3] if self.is_track else None
@property
@lru_cache(maxsize=2) # maxsize 1 should suffice
def xywh(self):
"""
Convert bounding boxes from [x1, y1, x2, y2] format to [x, y, width, height] format.
Returns:
(torch.Tensor | numpy.ndarray): Boxes in [x, y, width, height] format, where x, y are the coordinates of
the top-left corner of the bounding box, width, height are the dimensions of the bounding box and the
shape of the returned tensor is (N, 4), where N is the number of boxes.
Examples:
>>> boxes = Boxes(torch.tensor([[100, 50, 150, 100], [200, 150, 300, 250]]), orig_shape=(480, 640))
>>> xywh = boxes.xywh
>>> print(xywh)
tensor([[100.0000, 50.0000, 50.0000, 50.0000],
[200.0000, 150.0000, 100.0000, 100.0000]])
"""
# 将边界框从[x1, y1, x2, y2]格式转换为[x, y, width, height]格式
return ops.xyxy2xywh(self.xyxy)
@property
@lru_cache(maxsize=2)
def xyxyn(self):
"""
Returns normalized bounding box coordinates relative to the original image size.
This property calculates and returns the bounding box coordinates in [x1, y1, x2, y2] format,
normalized to the range [0, 1] based on the original image dimensions.
Returns:
(torch.Tensor | numpy.ndarray): Normalized bounding box coordinates with shape (N, 4), where N is
the number of boxes. Each row contains [x1, y1, x2, y2] values normalized to [0, 1].
Examples:
>>> boxes = Boxes(torch.tensor([[100, 50, 300, 400, 0.9, 0]]), orig_shape=(480, 640))
>>> normalized = boxes.xyxyn
>>> print(normalized)
tensor([[0.1562, 0.1042, 0.4688, 0.8333]])
"""
# Clone the bounding box tensor if it's a torch.Tensor; otherwise, create a numpy copy
xyxy = self.xyxy.clone() if isinstance(self.xyxy, torch.Tensor) else np.copy(self.xyxy)
# Normalize x1 and x2 coordinates by dividing with the width of the original image
xyxy[..., [0, 2]] /= self.orig_shape[1]
# Normalize y1 and y2 coordinates by dividing with the height of the original image
xyxy[..., [1, 3]] /= self.orig_shape[0]
# Return the normalized bounding box coordinates
return xyxy
@property
@lru_cache(maxsize=2)
def xywhn(self):
"""
Returns normalized bounding boxes in [x, y, width, height] format.
This property calculates and returns the normalized bounding box coordinates in the format
[x_center, y_center, width, height], where all values are relative to the original image dimensions.
Returns:
(torch.Tensor | numpy.ndarray): Normalized bounding boxes with shape (N, 4), where N is the
number of boxes. Each row contains [x_center, y_center, width, height] values normalized
to [0, 1] based on the original image dimensions.
Examples:
>>> boxes = Boxes(torch.tensor([[100, 50, 150, 100, 0.9, 0]]), orig_shape=(480, 640))
>>> normalized = boxes.xywhn
>>> print(normalized)
tensor([[0.1953, 0.1562, 0.0781, 0.1042]])
"""
# Convert bounding box coordinates from [x1, y1, x2, y2] to [x_center, y_center, width, height]
xywh = ops.xyxy2xywh(self.xyxy)
# Normalize x_center and width by dividing with the width of the original image
xywh[..., [0, 2]] /= self.orig_shape[1]
# Normalize y_center and height by dividing with the height of the original image
xywh[..., [1, 3]] /= self.orig_shape[0]
# Return the normalized bounding boxes in [x, y, width, height] format
return xywh
@property
@lru_cache(maxsize=1)
def xy(self) -> List[np.ndarray]:
"""
Property method that caches and returns a list of pixel coordinates of segmentation masks.
Returns:
List[np.ndarray]: A list where each element is a numpy array representing pixel coordinates
of a segmentation mask.
"""
return [self.data[i].nonzero()[:, :2] for i in range(self.data.shape[0])]
# 定义一个方法 xyn，用于返回分割掩模的归一化 xy 坐标
def xyn(self):
"""
Returns normalized xy-coordinates of the segmentation masks.
This property calculates and caches the normalized xy-coordinates of the segmentation masks. The coordinates
are normalized relative to the original image shape.
Returns:
(List[numpy.ndarray]): A list of numpy arrays, where each array contains the normalized xy-coordinates
of a single segmentation mask. Each array has shape (N, 2), where N is the number of points in the
mask contour.
Examples:
>>> results = model('image.jpg')
>>> masks = results[0].masks
>>> normalized_coords = masks.xyn
>>> print(normalized_coords[0]) # Normalized coordinates of the first mask
"""
# 使用 ops.masks2segments 方法将分割数据转换为分割段的列表，并对每个段的坐标进行归一化处理
return [
ops.scale_coords(self.data.shape[1:], x, self.orig_shape, normalize=True)
for x in ops.masks2segments(self.data)
]
@property
@lru_cache(maxsize=1)
# 使用 property 装饰器定义一个只读属性 xy，返回每个掩模张量中每个分割的像素坐标 [x, y]
def xy(self):
"""
Returns the [x, y] pixel coordinates for each segment in the mask tensor.
This property calculates and returns a list of pixel coordinates for each segmentation mask in the
Masks object. The coordinates are scaled to match the original image dimensions.
Returns:
(List[numpy.ndarray]): A list of numpy arrays, where each array contains the [x, y] pixel
coordinates for a single segmentation mask. Each array has shape (N, 2), where N is the
number of points in the segment.
Examples:
>>> results = model('image.jpg')
>>> masks = results[0].masks
>>> xy_coords = masks.xy
>>> print(len(xy_coords)) # Number of masks
>>> print(xy_coords[0].shape) # Shape of first mask's coordinates
"""
# 使用 ops.masks2segments 方法将分割数据转换为分割段的列表，并根据原始图像尺寸对每个段的坐标进行缩放处理
return [
ops.scale_coords(self.data.shape[1:], x, self.orig_shape, normalize=False)
for x in ops.masks2segments(self.data)
]
# 定义 Keypoints 类，继承自 BaseTensor
class Keypoints(BaseTensor):
"""
A class for storing and manipulating detection keypoints.
This class encapsulates functionality for handling keypoint data, including coordinate manipulation,
normalization, and confidence values.
Attributes:
data (torch.Tensor): The raw tensor containing keypoint data.
orig_shape (Tuple[int, int]): The original image dimensions (height, width).
has_visible (bool): Indicates whether visibility information is available for keypoints.
xy (torch.Tensor): Keypoint coordinates in [x, y] format.
xyn (torch.Tensor): Normalized keypoint coordinates in [x, y] format, relative to orig_shape.
conf (torch.Tensor): Confidence values for each keypoint, if available.
Methods:
cpu(): Returns a copy of the keypoints tensor on CPU memory.
numpy(): Returns a copy of the keypoints tensor as a numpy array.
cuda(): Returns a copy of the keypoints tensor on GPU memory.
to(*args, **kwargs): Returns a copy of the keypoints tensor with specified device and dtype.
Examples:
>>> import torch
>>> from ultralytics.engine.results import Keypoints
>>> keypoints_data = torch.rand(1, 17, 3) # 1 detection, 17 keypoints, (x, y, conf)
>>> orig_shape = (480, 640) # Original image shape (height, width)
>>> keypoints = Keypoints(keypoints_data, orig_shape)
>>> print(keypoints.xy.shape) # Access xy coordinates
>>> print(keypoints.conf) # Access confidence values
>>> keypoints_cpu = keypoints.cpu() # Move keypoints to CPU
"""
# 使用 smart_inference_mode 装饰器，用于避免在处理 keypoints 时出现小于 confidence 阈值的情况
@smart_inference_mode()
def __init__(self, keypoints, orig_shape) -> None:
"""
Initializes the Keypoints object with detection keypoints and original image dimensions.
This method processes the input keypoints tensor, handling both 2D and 3D formats. For 3D tensors
(x, y, confidence), it masks out low-confidence keypoints by setting their coordinates to zero.
Args:
keypoints (torch.Tensor): A tensor containing keypoint data. Shape can be either:
- (num_objects, num_keypoints, 2) for x, y coordinates only
- (num_objects, num_keypoints, 3) for x, y coordinates and confidence scores
orig_shape (Tuple[int, int]): The original image dimensions (height, width).
Examples:
>>> kpts = torch.rand(1, 17, 3) # 1 object, 17 keypoints (COCO format), x,y,conf
>>> orig_shape = (720, 1280) # Original image height, width
>>> keypoints = Keypoints(kpts, orig_shape)
"""
if keypoints.ndim == 2:
keypoints = keypoints[None, :] # 将2维关键点数据转换为3维（添加一个维度）
if keypoints.shape[2] == 3: # x, y, conf
mask = keypoints[..., 2] < 0.5 # 创建掩码，标记低置信度的关键点
keypoints[..., :2][mask] = 0 # 将低置信度关键点的坐标设置为0
super().__init__(keypoints, orig_shape) # 调用父类的初始化方法
self.has_visible = self.data.shape[-1] == 3 # 检查是否包含可见关键点信息
@property
@lru_cache(maxsize=1)
def xy(self):
"""
Returns x, y coordinates of keypoints.
Returns:
(torch.Tensor): A tensor containing the x, y coordinates of keypoints with shape (N, K, 2), where N is
the number of detections and K is the number of keypoints per detection.
Examples:
>>> results = model('image.jpg')
>>> keypoints = results[0].keypoints
>>> xy = keypoints.xy
>>> print(xy.shape) # (N, K, 2)
>>> print(xy[0]) # x, y coordinates of keypoints for first detection
Notes:
- The returned coordinates are in pixel units relative to the original image dimensions.
- If keypoints were initialized with confidence values, only keypoints with confidence >= 0.5 are returned.
- This property uses LRU caching to improve performance on repeated access.
"""
return self.data[..., :2] # 返回关键点的 x, y 坐标信息
@property
@lru_cache(maxsize=1)
这些注释解释了初始化方法和两个属性方法的功能及其作用，确保代码的每一部分都得到了清晰的解释。
def xyn(self):
"""
Returns normalized coordinates (x, y) of keypoints relative to the original image size.
Returns:
(torch.Tensor | numpy.ndarray): A tensor or array of shape (N, K, 2) containing normalized keypoint
coordinates, where N is the number of instances, K is the number of keypoints, and the last
dimension contains [x, y] values in the range [0, 1].
Examples:
>>> keypoints = Keypoints(torch.rand(1, 17, 2), orig_shape=(480, 640))
>>> normalized_kpts = keypoints.xyn
>>> print(normalized_kpts.shape)
torch.Size([1, 17, 2])
"""
# Clone the keypoint coordinates if they are a torch.Tensor; otherwise, create a numpy copy
xy = self.xy.clone() if isinstance(self.xy, torch.Tensor) else np.copy(self.xy)
# Normalize x coordinates by dividing by the width of the original image
xy[..., 0] /= self.orig_shape[1]
# Normalize y coordinates by dividing by the height of the original image
xy[..., 1] /= self.orig_shape[0]
# Return the normalized keypoint coordinates
return xy
@property
@lru_cache(maxsize=1)
def conf(self):
"""
Returns confidence values for each keypoint.
Returns:
(torch.Tensor | None): A tensor containing confidence scores for each keypoint if available,
otherwise None. Shape is (num_detections, num_keypoints) for batched data or (num_keypoints,)
for single detection.
Examples:
>>> keypoints = Keypoints(torch.rand(1, 17, 3), orig_shape=(640, 640)) # 1 detection, 17 keypoints
>>> conf = keypoints.conf
>>> print(conf.shape) # torch.Size([1, 17])
"""
# Return confidence scores if keypoints have visibility information; otherwise return None
return self.data[..., 2] if self.has_visible else None
@property
@lru_cache(maxsize=1)
def top1(self):
"""
Return the index of the class with the highest probability.
This property computes and returns the index of the class with the highest probability
from the stored classification probabilities.
Returns:
int: Index of the class with the highest probability.
"""
def top1(self):
"""
Returns the index of the class with the highest probability.
Returns:
(int): Index of the class with the highest probability.
Examples:
>>> probs = Probs(torch.tensor([0.1, 0.3, 0.6]))
>>> probs.top1
2
"""
# 返回数据张量中最大值的索引，即最高概率对应的类别索引
return int(self.data.argmax())
@property
@lru_cache(maxsize=1)
def top5(self):
"""
Returns the indices of the top 5 class probabilities.
Returns:
(List[int]): A list containing the indices of the top 5 class probabilities, sorted in descending order.
Examples:
>>> probs = Probs(torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5]))
>>> print(probs.top5)
[4, 3, 2, 1, 0]
"""
# 返回数据张量中前五个最大值的索引列表，按降序排列
return (-self.data).argsort(0)[:5].tolist() # this way works with both torch and numpy.
@property
@lru_cache(maxsize=1)
def top1conf(self):
"""
Returns the confidence score of the highest probability class.
This property retrieves the confidence score (probability) of the class with the highest predicted probability
from the classification results.
Returns:
(torch.Tensor | numpy.ndarray): A tensor containing the confidence score of the top 1 class.
Examples:
>>> results = model('image.jpg') # classify an image
>>> probs = results[0].probs # get classification probabilities
>>> top1_confidence = probs.top1conf # get confidence of top 1 class
>>> print(f"Top 1 class confidence: {top1_confidence.item():.4f}")
"""
# 返回数据张量中最高概率类别的置信度分数
return self.data[self.top1]
@property
@lru_cache(maxsize=1)
def top5conf(self):
"""
Returns confidence scores for the top 5 classification predictions.
This property retrieves the confidence scores corresponding to the top 5 class probabilities
predicted by the model. It provides a quick way to access the most likely class predictions
along with their associated confidence levels.
Returns:
(torch.Tensor | numpy.ndarray): A tensor or array containing the confidence scores for the
top 5 predicted classes, sorted in descending order of probability.
Examples:
>>> results = model('image.jpg')
>>> probs = results[0].probs
>>> top5_conf = probs.top5conf
>>> print(top5_conf) # Prints confidence scores for top 5 classes
"""
# 返回数据张量中前五个最大概率类别的置信度分数
return self.data[self.top5]
# 定义一个名为 OBB 的类，继承自 BaseTensor
class OBB(BaseTensor):
"""
A class for storing and manipulating Oriented Bounding Boxes (OBB).
This class provides functionality to handle oriented bounding boxes, including conversion between
different formats, normalization, and access to various properties of the boxes.
Attributes:
data (torch.Tensor): The raw OBB tensor containing box coordinates and associated data.
orig_shape (tuple): Original image size as (height, width).
is_track (bool): Indicates whether tracking IDs are included in the box data.
xywhr (torch.Tensor | numpy.ndarray): Boxes in [x_center, y_center, width, height, rotation] format.
conf (torch.Tensor | numpy.ndarray): Confidence scores for each box.
cls (torch.Tensor | numpy.ndarray): Class labels for each box.
id (torch.Tensor | numpy.ndarray): Tracking IDs for each box, if available.
xyxyxyxy (torch.Tensor | numpy.ndarray): Boxes in 8-point [x1, y1, x2, y2, x3, y3, x4, y4] format.
xyxyxyxyn (torch.Tensor | numpy.ndarray): Normalized 8-point coordinates relative to orig_shape.
xyxy (torch.Tensor | numpy.ndarray): Axis-aligned bounding boxes in [x1, y1, x2, y2] format.
Methods:
cpu(): Returns a copy of the OBB object with all tensors on CPU memory.
numpy(): Returns a copy of the OBB object with all tensors as numpy arrays.
cuda(): Returns a copy of the OBB object with all tensors on GPU memory.
to(*args, **kwargs): Returns a copy of the OBB object with tensors on specified device and dtype.
Examples:
>>> boxes = torch.tensor([[100, 50, 150, 100, 30, 0.9, 0]]) # xywhr, conf, cls
>>> obb = OBB(boxes, orig_shape=(480, 640))
>>> print(obb.xyxyxyxy)
>>> print(obb.conf)
>>> print(obb.cls)
"""
def __init__(self, boxes, orig_shape) -> None:
"""
Initialize an OBB (Oriented Bounding Box) instance with oriented bounding box data and original image shape.
This class stores and manipulates Oriented Bounding Boxes (OBB) for object detection tasks. It provides
various properties and methods to access and transform the OBB data.
Args:
boxes (torch.Tensor | numpy.ndarray): A tensor or numpy array containing the detection boxes,
with shape (num_boxes, 7) or (num_boxes, 8). The last two columns contain confidence and class values.
If present, the third last column contains track IDs, and the fifth column contains rotation.
orig_shape (Tuple[int, int]): Original image size, in the format (height, width).
Attributes:
data (torch.Tensor | numpy.ndarray): The raw OBB tensor.
orig_shape (Tuple[int, int]): The original image shape.
is_track (bool): Whether the boxes include tracking IDs.
Raises:
AssertionError: If the number of values per box is not 7 or 8.
Examples:
>>> import torch
>>> boxes = torch.rand(3, 7) # 3 boxes with 7 values each
>>> orig_shape = (640, 480)
>>> obb = OBB(boxes, orig_shape)
>>> print(obb.xywhr) # Access the boxes in xywhr format
"""
if boxes.ndim == 1:
boxes = boxes[None, :] # 将一维的boxes转换为二维，以确保正确的形状
n = boxes.shape[-1] # 获取最后一个维度的大小，即每个box的值个数
assert n in {7, 8}, f"expected 7 or 8 values but got {n}" # 断言检查每个box的值个数是否为7或8
super().__init__(boxes, orig_shape) # 调用父类的初始化方法，传入boxes和orig_shape
self.is_track = n == 8 # 判断是否包含track IDs，如果值个数为8，则包含
self.orig_shape = orig_shape # 存储原始图像的形状信息
@property
def xywhr(self):
"""
Returns boxes in [x_center, y_center, width, height, rotation] format.
Returns:
(torch.Tensor | numpy.ndarray): A tensor or numpy array containing the oriented bounding boxes with format
[x_center, y_center, width, height, rotation]. The shape is (N, 5) where N is the number of boxes.
Examples:
>>> results = model('image.jpg')
>>> obb = results[0].obb
>>> xywhr = obb.xywhr
>>> print(xywhr.shape)
torch.Size([3, 5])
"""
return self.data[:, :5] # 返回包含 [x_center, y_center, width, height, rotation] 的部分数据
def conf(self):
"""
Returns the confidence scores for Oriented Bounding Boxes (OBBs).
This property retrieves the confidence values associated with each OBB detection. The confidence score
represents the model's certainty in the detection.
Returns:
(torch.Tensor | numpy.ndarray): A tensor or numpy array of shape (N,) containing confidence scores
for N detections, where each score is in the range [0, 1].
Examples:
>>> results = model('image.jpg')
>>> obb_result = results[0].obb
>>> confidence_scores = obb_result.conf
>>> print(confidence_scores)
"""
# 返回包含所有检测结果置信度的数据列
return self.data[:, -2]
@property
def cls(self):
"""
Returns the class values of the oriented bounding boxes.
Returns:
(torch.Tensor | numpy.ndarray): A tensor or numpy array containing the class values for each oriented
bounding box. The shape is (N,), where N is the number of boxes.
Examples:
>>> results = model('image.jpg')
>>> result = results[0]
>>> obb = result.obb
>>> class_values = obb.cls
>>> print(class_values)
"""
# 返回包含所有检测结果类别的数据列
return self.data[:, -1]
@property
def id(self):
"""
Returns the tracking IDs of the oriented bounding boxes (if available).
Returns:
(torch.Tensor | numpy.ndarray | None): A tensor or numpy array containing the tracking IDs for each
oriented bounding box. Returns None if tracking IDs are not available.
Examples:
>>> results = model('image.jpg', tracker=True) # Run inference with tracking
>>> for result in results:
... if result.obb is not None:
... track_ids = result.obb.id
... if track_ids is not None:
... print(f"Tracking IDs: {track_ids}")
"""
# 如果追踪信息可用，则返回包含所有检测结果追踪 ID 的数据列；否则返回 None
return self.data[:, -3] if self.is_track else None
@property
@lru_cache(maxsize=2)
def xyxyxyxy(self):
"""
Converts OBB format to 8-point (xyxyxyxy) coordinate format for rotated bounding boxes.
Returns:
(torch.Tensor | numpy.ndarray): Rotated bounding boxes in xyxyxyxy format with shape (N, 4, 2), where N is
the number of boxes. Each box is represented by 4 points (x, y), starting from the top-left corner and
moving clockwise.
Examples:
>>> obb = OBB(torch.tensor([[100, 100, 50, 30, 0.5, 0.9, 0]]), orig_shape=(640, 640))
>>> xyxyxyxy = obb.xyxyxyxy
>>> print(xyxyxyxy.shape)
torch.Size([1, 4, 2])
"""
# 将 OBB 格式转换为 xyxyxyxy 格式的旋转边界框坐标
return ops.xywhr2xyxyxyxy(self.xywhr)
@property
@lru_cache(maxsize=2)
def xyxyxyxyn(self):
"""
Converts rotated bounding boxes to normalized xyxyxyxy format.
Returns:
(torch.Tensor | numpy.ndarray): Normalized rotated bounding boxes in xyxyxyxy format with shape (N, 4, 2),
where N is the number of boxes. Each box is represented by 4 points (x, y), normalized relative to
the original image dimensions.
Examples:
>>> obb = OBB(torch.rand(10, 7), orig_shape=(640, 480)) # 10 random OBBs
>>> normalized_boxes = obb.xyxyxyxyn
>>> print(normalized_boxes.shape)
torch.Size([10, 4, 2])
"""
# 创建副本以确保不修改原始数据
xyxyxyxyn = self.xyxyxyxy.clone() if isinstance(self.xyxyxyxy, torch.Tensor) else np.copy(self.xyxyxyxy)
# 归一化 x 坐标
xyxyxyxyn[..., 0] /= self.orig_shape[1]
# 归一化 y 坐标
xyxyxyxyn[..., 1] /= self.orig_shape[0]
# 返回归一化后的坐标
return xyxyxyxyn
@property
@lru_cache(maxsize=2)
def xyxy(self):
"""
Converts oriented bounding boxes (OBB) to axis-aligned bounding boxes in xyxy format.
This property calculates the minimal enclosing rectangle for each oriented bounding box and returns it in
xyxy format (x1, y1, x2, y2). This is useful for operations that require axis-aligned bounding boxes, such
as IoU calculation with non-rotated boxes.
Returns:
(torch.Tensor | numpy.ndarray): Axis-aligned bounding boxes in xyxy format with shape (N, 4), where N
is the number of boxes. Each row contains [x1, y1, x2, y2] coordinates.
Examples:
>>> import torch
>>> from ultralytics import YOLO
>>> model = YOLO('yolov8n-obb.pt')
>>> results = model('path/to/image.jpg')
>>> for result in results:
... obb = result.obb
... if obb is not None:
... xyxy_boxes = obb.xyxy
... print(xyxy_boxes.shape) # (N, 4)
Notes:
- This method approximates the OBB by its minimal enclosing rectangle.
- The returned format is compatible with standard object detection metrics and visualization tools.
- The property uses caching to improve performance for repeated access.
"""
# 提取 x 和 y 坐标
x = self.xyxyxyxy[..., 0]
y = self.xyxyxyxy[..., 1]
# 根据最小值和最大值创建包围框
return (
torch.stack([x.amin(1), y.amin(1), x.amax(1), y.amax(1)], -1)
if isinstance(x, torch.Tensor)
else np.stack([x.min(1), y.min(1), x.max(1), y.max(1)], -1)
)
.\yolov8\ultralytics\engine\trainer.py
# Ultralytics YOLO 🚀, AGPL-3.0 license
"""
Train a model on a dataset.
Usage:
$ yolo mode=train model=yolov8n.pt data=coco8.yaml imgsz=640 epochs=100 batch=16
"""
# 导入需要的库和模块
import gc # 垃圾回收模块
import math # 数学计算模块
import os # 系统操作模块
import subprocess # 子进程管理模块
import time # 时间操作模块
import warnings # 警告处理模块
from copy import deepcopy # 深拷贝函数
from datetime import datetime, timedelta # 日期时间操作模块
from pathlib import Path # 路径操作模块
import numpy as np # 数组操作模块
import torch # PyTorch深度学习库
from torch import distributed as dist # 分布式训练模块
from torch import nn, optim # 神经网络模块，优化器模块
# 导入Ultralytics库中的各种功能和工具
from ultralytics.cfg import get_cfg, get_save_dir # 获取配置和保存目录
from ultralytics.data.utils import check_cls_dataset, check_det_dataset # 检查分类和检测数据集的实用工具
from ultralytics.nn.tasks import attempt_load_one_weight, attempt_load_weights # 加载模型权重的函数
from ultralytics.utils import (
DEFAULT_CFG, # 默认配置
LOGGER, # 日志记录器
RANK, # 排名
TQDM, # 进度条
__version__, # Ultralytics版本号
callbacks, # 回调函数
clean_url, # 清理URL
colorstr, # 字符串颜色处理
emojis, # 表情符号
yaml_save, # 保存YAML文件
)
from ultralytics.utils.autobatch import check_train_batch_size # 检查训练批次大小
from ultralytics.utils.checks import (
check_amp, # 检查AMP（自动混合精度）支持
check_file, # 检查文件
check_imgsz, # 检查图像尺寸
check_model_file_from_stem, # 从stem检查模型文件
print_args, # 打印参数
)
from ultralytics.utils.dist import ddp_cleanup, generate_ddp_command # 分布式数据并行清理和命令生成
from ultralytics.utils.files import get_latest_run # 获取最新运行文件
from ultralytics.utils.torch_utils import (
EarlyStopping, # 早停
ModelEMA, # 模型指数移动平均
autocast, # 自动类型转换
convert_optimizer_state_dict_to_fp16, # 将优化器状态转换为FP16
init_seeds, # 初始化随机种子
one_cycle, # 单循环学习率策略
select_device, # 选择设备
strip_optimizer, # 剥离优化器
torch_distributed_zero_first, # 分布式训练的首次清零
)
# 定义基础训练器类
class BaseTrainer:
"""
BaseTrainer.
A base class for creating trainers.
Attributes:
args (SimpleNamespace): Configuration for the trainer.
validator (BaseValidator): Validator instance.
model (nn.Module): Model instance.
callbacks (defaultdict): Dictionary of callbacks.
save_dir (Path): Directory to save results.
wdir (Path): Directory to save weights.
last (Path): Path to the last checkpoint.
best (Path): Path to the best checkpoint.
save_period (int): Save checkpoint every x epochs (disabled if < 1).
batch_size (int): Batch size for training.
epochs (int): Number of epochs to train for.
start_epoch (int): Starting epoch for training.
device (torch.device): Device to use for training.
amp (bool): Flag to enable AMP (Automatic Mixed Precision).
scaler (amp.GradScaler): Gradient scaler for AMP.
data (str): Path to data.
trainset (torch.utils.data.Dataset): Training dataset.
testset (torch.utils.data.Dataset): Testing dataset.
ema (nn.Module): EMA (Exponential Moving Average) of the model.
resume (bool): Resume training from a checkpoint.
lf (nn.Module): Loss function.
scheduler (torch.optim.lr_scheduler._LRScheduler): Learning rate scheduler.
best_fitness (float): The best fitness value achieved.
fitness (float): Current fitness value.
loss (float): Current loss value.
tloss (float): Total loss value.
loss_names (list): List of loss names.
csv (Path): Path to results CSV file.
"""
# 初始化 BaseTrainer 类的构造函数
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
"""
Initializes the BaseTrainer class.
Args:
cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CFG.
overrides (dict, optional): Configuration overrides. Defaults to None.
_callbacks (list, optional): List of callbacks. Defaults to None.
"""
# 获取配置参数
self.args = get_cfg(cfg, overrides)
# 检查是否需要恢复训练状态
self.check_resume(overrides)
# 选择设备
self.device = select_device(self.args.device, self.args.batch)
self.validator = None # 初始化验证器为 None
self.metrics = None # 初始化指标为 None
self.plots = {} # 初始化图形字典为空
# 初始化随机种子
init_seeds(self.args.seed + 1 + RANK, deterministic=self.args.deterministic)
# 目录设置
self.save_dir = get_save_dir(self.args) # 获取保存目录
self.args.name = self.save_dir.name # 更新日志记录器的名称
self.wdir = self.save_dir / "weights" # 权重保存目录
if RANK in {-1, 0}:
self.wdir.mkdir(parents=True, exist_ok=True) # 创建目录
self.args.save_dir = str(self.save_dir) # 设置保存目录路径
yaml_save(self.save_dir / "args.yaml", vars(self.args)) # 保存运行参数到 YAML 文件
self.last, self.best = self.wdir / "last.pt", self.wdir / "best.pt" # 检查点文件路径
self.save_period = self.args.save_period # 保存周期设置
self.batch_size = self.args.batch # 批量大小
self.epochs = self.args.epochs # 训练周期数
self.start_epoch = 0 # 起始周期
if RANK == -1:
print_args(vars(self.args)) # 打印参数信息
# 设备设置
if self.device.type in {"cpu", "mps"}:
self.args.workers = 0 # 若使用 CPU 或者 mps 设备，设置 workers 为 0，提升训练速度
# 模型和数据集
self.model = check_model_file_from_stem(self.args.model) # 检查模型文件，添加后缀（例如 yolov8n -> yolov8n.pt）
with torch_distributed_zero_first(RANK): # 避免多次自动下载数据集
self.trainset, self.testset = self.get_dataset() # 获取训练集和测试集
self.ema = None # 指数移动平均参数初始化为 None
# 优化器工具初始化
self.lf = None # 损失函数初始化为 None
self.scheduler = None # 调度器初始化为 None
# 每个周期的指标
self.best_fitness = None # 最佳适应度初始化为 None
self.fitness = None # 适应度初始化为 None
self.loss = None # 损失初始化为 None
self.tloss = None # 总损失初始化为 None
self.loss_names = ["Loss"] # 损失名称列表
self.csv = self.save_dir / "results.csv" # CSV 文件路径
self.plot_idx = [0, 1, 2] # 绘图索引
# HUB
self.hub_session = None # HUB 会话初始化为 None
# 回调函数
self.callbacks = _callbacks or callbacks.get_default_callbacks() # 获取默认或自定义回调函数列表
if RANK in {-1, 0}:
callbacks.add_integration_callbacks(self) # 添加集成回调函数
def add_callback(self, event: str, callback):
"""Appends the given callback to the event."""
self.callbacks[event].append(callback) # 将给定的回调函数添加到事件对应的回调函数列表
def set_callback(self, event: str, callback):
"""Sets the given callback for the event, overriding any existing callbacks."""
self.callbacks[event] = [callback] # 设置给定事件的回调函数，覆盖现有的所有回调函数
def run_callbacks(self, event: str):
"""Runs all existing callbacks associated with the given event."""
for callback in self.callbacks.get(event, []): # 遍历给定事件关联的所有回调函数
callback(self) # 执行回调函数
def train(self):
"""
Allow device='', device=None on Multi-GPU systems to default to device=0.
"""
# Determine the number of GPUs to be used based on the provided device configuration
if isinstance(self.args.device, str) and len(self.args.device): # i.e. device='0' or device='0,1,2,3'
world_size = len(self.args.device.split(","))
elif isinstance(self.args.device, (tuple, list)): # i.e. device=[0, 1, 2, 3] (multi-GPU from CLI is list)
world_size = len(self.args.device)
elif torch.cuda.is_available(): # i.e. device=None or device='' or device=number
world_size = 1 # default to device 0
else: # i.e. device='cpu' or 'mps'
world_size = 0
# Run subprocess if DDP (Distributed Data Parallel) training is enabled and not in a subprocess
if world_size > 1 and "LOCAL_RANK" not in os.environ:
# Argument checks and adjustments for Multi-GPU training
if self.args.rect:
LOGGER.warning("WARNING ⚠️ 'rect=True' is incompatible with Multi-GPU training, setting 'rect=False'")
self.args.rect = False
if self.args.batch < 1.0:
LOGGER.warning(
"WARNING ⚠️ 'batch<1' for AutoBatch is incompatible with Multi-GPU training, setting "
"default 'batch=16'"
)
self.args.batch = 16
# Generate and execute the DDP command
cmd, file = generate_ddp_command(world_size, self)
try:
LOGGER.info(f'{colorstr("DDP:")} debug command {" ".join(cmd)}')
subprocess.run(cmd, check=True)
except Exception as e:
raise e
finally:
# Clean up DDP resources
ddp_cleanup(self, str(file))
else:
# Perform single-GPU or CPU training
self._do_train(world_size)
def _setup_scheduler(self):
"""
Initialize training learning rate scheduler based on arguments.
"""
if self.args.cos_lr:
# Use cosine learning rate schedule
self.lf = one_cycle(1, self.args.lrf, self.epochs) # cosine 1->hyp['lrf']
else:
# Use linear learning rate schedule
self.lf = lambda x: max(1 - x / self.epochs, 0) * (1.0 - self.args.lrf) + self.args.lrf # linear
# Initialize LambdaLR scheduler with the defined learning rate function
self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf)
def _setup_ddp(self, world_size):
"""
Initializes and configures DistributedDataParallel (DDP) for training.
"""
# Set CUDA device for the current process rank
torch.cuda.set_device(RANK)
self.device = torch.device("cuda", RANK)
# Configure environment for DDP training
os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1" # set to enforce timeout
# Initialize the process group for distributed training
dist.init_process_group(
backend="nccl" if dist.is_nccl_available() else "gloo",
timeout=timedelta(seconds=10800), # 3 hours
rank=RANK,
world_size=world_size,
)
def save_model(self):
"""
Save model training checkpoints with additional metadata.
"""
import io # 导入io模块，用于处理字节流操作
import pandas as pd # 导入pandas模块，用于处理CSV文件和数据结构
# Serialize ckpt to a byte buffer once (faster than repeated torch.save() calls)
buffer = io.BytesIO() # 创建一个字节流缓冲区对象
torch.save(
{
"epoch": self.epoch, # 保存当前训练轮次
"best_fitness": self.best_fitness, # 保存最佳训练效果指标
"model": None, # 模型信息（由EMA派生的恢复和最终检查点）
"ema": deepcopy(self.ema.ema).half(), # 深度复制EMA对象的EMA部分，并转换为半精度
"updates": self.ema.updates, # 保存EMA对象的更新次数
"optimizer": convert_optimizer_state_dict_to_fp16(deepcopy(self.optimizer.state_dict())), # 转换优化器状态字典为半精度
"train_args": vars(self.args), # 保存训练参数的字典形式
"train_metrics": {**self.metrics, **{"fitness": self.fitness}}, # 保存训练指标和适应度指标
"train_results": {k.strip(): v for k, v in pd.read_csv(self.csv).to_dict(orient="list").items()}, # 读取CSV文件并将结果保存为字典
"date": datetime.now().isoformat(), # 保存当前时间的ISO格式字符串
"version": __version__, # 保存当前版本号
"license": "AGPL-3.0 (https://ultralytics.com/license)", # 许可证信息
"docs": "https://docs.ultralytics.com", # 文档链接
},
buffer,
)
serialized_ckpt = buffer.getvalue() # 获取序列化后的检查点内容以保存
# Save checkpoints
self.last.write_bytes(serialized_ckpt) # 保存最后的检查点文件（last.pt）
if self.best_fitness == self.fitness:
self.best.write_bytes(serialized_ckpt) # 若当前适应度等于最佳适应度，则保存最佳检查点文件（best.pt）
if (self.save_period > 0) and (self.epoch % self.save_period == 0):
(self.wdir / f"epoch{self.epoch}.pt").write_bytes(serialized_ckpt) # 根据当前轮次保存检查点文件，如'epoch3.pt'
def get_dataset(self):
"""
Get train, val path from data dict if it exists.
Returns None if data format is not recognized.
"""
try:
if self.args.task == "classify": # 如果任务为分类
data = check_cls_dataset(self.args.data) # 检查并获取分类数据集
elif self.args.data.split(".")[-1] in {"yaml", "yml"} or self.args.task in {
"detect",
"segment",
"pose",
"obb",
}: # 如果数据文件格式为yaml或yml，或者任务是检测、分割、姿态或obb
data = check_det_dataset(self.args.data) # 检查并获取检测数据集
if "yaml_file" in data:
self.args.data = data["yaml_file"] # 更新self.args.data以验证'yolo train data=url.zip'的使用
except Exception as e:
raise RuntimeError(emojis(f"Dataset '{clean_url(self.args.data)}' error ❌ {e}")) from e # 抛出数据集错误异常
self.data = data # 将获取的数据集保存到self.data
return data["train"], data.get("val") or data.get("test") # 返回训练集路径，如果不存在则返回验证集或测试集路径
def setup_model(self):
"""
Load/create/download model for any task.
Checks if a model is already loaded. If not, attempts to load weights from a '.pt' file or a specified pretrained path,
then initializes the model using the loaded configuration and weights.
"""
if isinstance(self.model, torch.nn.Module): # if model is loaded beforehand. No setup needed
return
cfg, weights = self.model, None
ckpt = None
if str(self.model).endswith(".pt"):
# Attempt to load weights and configuration from a '.pt' file
weights, ckpt = attempt_load_one_weight(self.model)
cfg = weights.yaml
elif isinstance(self.args.pretrained, (str, Path)):
# Attempt to load weights from a specified pretrained path
weights, _ = attempt_load_one_weight(self.args.pretrained)
# Initialize the model using the retrieved configuration and weights
self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK == -1) # calls Model(cfg, weights)
return ckpt
def optimizer_step(self):
"""
Perform a single step of the training optimizer with gradient clipping and EMA update.
Unscales gradients, clips gradients to a maximum norm of 10.0, performs optimizer step, updates scaler,
and zeroes the gradients of the optimizer. If an Exponential Moving Average (EMA) is enabled, updates the EMA.
"""
self.scaler.unscale_(self.optimizer) # unscale gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0) # clip gradients
self.scaler.step(self.optimizer)
self.scaler.update()
self.optimizer.zero_grad()
if self.ema:
self.ema.update(self.model)
def preprocess_batch(self, batch):
"""
Allows custom preprocessing model inputs and ground truths depending on task type.
This function can be overridden to implement custom preprocessing logic for batches of data before feeding them
into the model.
"""
return batch
def validate(self):
"""
Runs validation on test set using self.validator.
Returns metrics and fitness. Updates the best fitness if a new best fitness is found during validation.
"""
metrics = self.validator(self)
fitness = metrics.pop("fitness", -self.loss.detach().cpu().numpy()) # use loss as fitness measure if not found
if not self.best_fitness or self.best_fitness < fitness:
self.best_fitness = fitness
return metrics, fitness
def get_model(self, cfg=None, weights=None, verbose=True):
"""
Get model and raise NotImplementedError for loading cfg files.
Raises NotImplementedError indicating that loading configuration files (cfg) is not supported by this method.
"""
raise NotImplementedError("This task trainer doesn't support loading cfg files")
def get_validator(self):
"""
Returns a NotImplementedError when the get_validator function is called.
Raises NotImplementedError indicating that the get_validator function is not implemented in this trainer.
"""
raise NotImplementedError("get_validator function not implemented in trainer")
def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
"""
Returns dataloader derived from torch.data.Dataloader.
Raises NotImplementedError indicating that the get_dataloader function is not implemented in this trainer.
"""
raise NotImplementedError("get_dataloader function not implemented in trainer")
def build_dataset(self, img_path, mode="train", batch=None):
"""
Build dataset.
Raises NotImplementedError indicating that the build_dataset function is not implemented in this trainer.
"""
raise NotImplementedError("build_dataset function not implemented in trainer")
def label_loss_items(self, loss_items=None, prefix="train"):
"""
Returns a loss dict with labelled training loss items tensor.
This function provides a labeled dictionary of loss items. It's particularly useful for tasks like segmentation and detection.
"""
return {"loss": loss_items} if loss_items is not None else ["loss"]
# 设置或更新模型训练前的参数
def set_model_attributes(self):
"""To set or update model parameters before training."""
self.model.names = self.data["names"]
# 构建用于训练 YOLO 模型的目标张量
def build_targets(self, preds, targets):
"""Builds target tensors for training YOLO model."""
pass
# 返回描述训练进度的字符串
def progress_string(self):
"""Returns a string describing training progress."""
return ""
# TODO: 可能需要将以下函数放入回调函数中
# 绘制 YOLO 训练过程中的训练样本图像
def plot_training_samples(self, batch, ni):
"""Plots training samples during YOLO training."""
pass
# 绘制 YOLO 模型的训练标签
def plot_training_labels(self):
"""Plots training labels for YOLO model."""
pass
# 将训练指标保存到 CSV 文件中
def save_metrics(self, metrics):
"""Saves training metrics to a CSV file."""
keys, vals = list(metrics.keys()), list(metrics.values())
n = len(metrics) + 1 # number of cols
s = "" if self.csv.exists() else (("%23s," * n % tuple(["epoch"] + keys)).rstrip(",") + "\n") # header
with open(self.csv, "a") as f:
f.write(s + ("%23.5g," * n % tuple([self.epoch + 1] + vals)).rstrip(",") + "\n")
# 绘制并可视化训练指标
def plot_metrics(self):
"""Plot and display metrics visually."""
pass
# 注册绘图信息（例如供回调函数使用）
def on_plot(self, name, data=None):
"""Registers plots (e.g. to be consumed in callbacks)"""
path = Path(name)
self.plots[path] = {"data": data, "timestamp": time.time()}
# 执行物体检测 YOLO 模型的最终评估和验证
def final_eval(self):
"""Performs final evaluation and validation for object detection YOLO model."""
for f in self.last, self.best:
if f.exists():
strip_optimizer(f) # 去除优化器信息
if f is self.best:
LOGGER.info(f"\nValidating {f}...")
self.validator.args.plots = self.args.plots
self.metrics = self.validator(model=f)
self.metrics.pop("fitness", None)
self.run_callbacks("on_fit_epoch_end")
# 检查是否存在恢复检查点，并根据需要更新参数
def check_resume(self, overrides):
# 获取恢复路径
resume = self.args.resume
# 如果有恢复路径
if resume:
try:
# 检查恢复路径是否是字符串或路径对象并且存在
exists = isinstance(resume, (str, Path)) and Path(resume).exists()
# 获取最新的运行检查点路径
last = Path(check_file(resume) if exists else get_latest_run())
# 检查恢复的数据 YAML 文件是否存在，否则强制重新下载数据集
ckpt_args = attempt_load_weights(last).args
if not Path(ckpt_args["data"]).exists():
ckpt_args["data"] = self.args.data
# 标记为需要恢复
resume = True
# 获取配置对象，并更新参数
self.args = get_cfg(ckpt_args)
self.args.model = self.args.resume = str(last) # 恢复模型路径
# 遍历可能需要更新的参数，例如减少内存占用或更新设备
for k in "imgsz", "batch", "device":
if k in overrides:
setattr(self.args, k, overrides[k])
except Exception as e:
# 如果出现异常，抛出文件未找到的错误
raise FileNotFoundError(
"Resume checkpoint not found. Please pass a valid checkpoint to resume from, "
"i.e. 'yolo train resume model=path/to/last.pt'"
) from e
# 将恢复状态更新到对象中
self.resume = resume
# 从给定的检查点恢复 YOLO 训练
def resume_training(self, ckpt):
# 如果检查点为空或者没有设置恢复标志，则直接返回
if ckpt is None or not self.resume:
return
# 初始化最佳适应度和开始的 epoch
best_fitness = 0.0
start_epoch = ckpt.get("epoch", -1) + 1
# 如果检查点中包含优化器状态，则加载优化器状态
if ckpt.get("optimizer", None) is not None:
self.optimizer.load_state_dict(ckpt["optimizer"]) # 加载优化器状态
best_fitness = ckpt["best_fitness"]
# 如果存在 EMA 并且检查点中包含 EMA 数据，则加载 EMA 状态
if self.ema and ckpt.get("ema"):
self.ema.ema.load_state_dict(ckpt["ema"].float().state_dict()) # 加载 EMA 状态
self.ema.updates = ckpt["updates"]
# 确保开始的 epoch 大于 0，否则输出错误信息
assert start_epoch > 0, (
f"{self.args.model} training to {self.epochs} epochs is finished, nothing to resume.\n"
f"Start a new training without resuming, i.e. 'yolo train model={self.args.model}'"
)
# 记录恢复训练的信息
LOGGER.info(f"Resuming training {self.args.model} from epoch {start_epoch + 1} to {self.epochs} total epochs")
# 如果总 epoch 小于开始的 epoch，则输出调试信息进行微调
if self.epochs < start_epoch:
LOGGER.info(
f"{self.model} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {self.epochs} more epochs."
)
self.epochs += ckpt["epoch"] # 进行额外的微调 epoch
# 更新最佳适应度和开始的 epoch
self.best_fitness = best_fitness
self.start_epoch = start_epoch
# 如果开始的 epoch 大于总 epoch 减去关闭混合图块的数量，则关闭数据加载器中的混合图块
if start_epoch > (self.epochs - self.args.close_mosaic):
self._close_dataloader_mosaic()
# 定义一个方法，用于关闭数据加载器的马赛克增强功能
def _close_dataloader_mosaic(self):
"""Update dataloaders to stop using mosaic augmentation."""
# 检查训练数据加载器的数据集是否具有 'mosaic' 属性
if hasattr(self.train_loader.dataset, "mosaic"):
# 将数据集的 'mosaic' 属性设置为 False，停止使用马赛克增强
self.train_loader.dataset.mosaic = False
# 再次检查训练数据加载器的数据集是否具有 'close_mosaic' 方法
if hasattr(self.train_loader.dataset, "close_mosaic"):
# 记录信息到日志，指示正在关闭数据加载器的马赛克增强
LOGGER.info("Closing dataloader mosaic")
# 调用数据集的 'close_mosaic' 方法，关闭马赛克增强，并传递额外的参数 'hyp=self.args'
self.train_loader.dataset.close_mosaic(hyp=self.args)
