关于.pt文件的使用

Question

我刚接触pytorch，现在已经使用yolov5训练好模型了。怎样使用best.pt文件输出测试图片

Answer 1

在使用 YOLOv5 训练好的模型进行测试时，可以使用以下步骤：

加载训练好的模型。
您可以使用 torch.load 函数加载模型文件 best.pt。例如：

import torch

model = torch.load('best.pt')

加载测试图片。
使用 Python 的图像处理库（如 Pillow、OpenCV 等）加载测试图片。例如：

from PIL import Image

image = Image.open('test.jpg')

对测试图片进行预处理。
在使用模型进行预测之前，需要将图片进行预处理，以符合模型的输入要求。例如，YOLOv5 模型的输入通常是一个 3 维的张量，形状为 (3, 416, 416)。因此，可以将图片转化为这种形状，并将像素值转化为浮点数。

import numpy as np

# 将图片转化为 numpy 数组
image_array = np.array(image)

# 将图片的像素值转化为浮点数
image_array = image_array.astype(np.float32)

# 将图片的像素值进行归一化
image_array /= 255.0

# 将图片的形状转化为 (3, 416, 416)
image_array = np.transpose(image_array, (2, 0, 1))
image_array = np.expand_dims(image_array, axis=0)

将图片输入模型并进行预测。
使用 model.eval() 将模型转化为评估模式，然后使用 model(image_array) 进行预测。例如：

model.eval()

prediction = model(image_array)

处理模型的输出并可视化。
模型的输出通常是一个列表，其中包含检测到的物体的位置、类别、置信度等信息。可以使用 Numpy、Matplotlib 等库处理这些信息，然后在图片上绘制出这些信息。例如：

# 处理模型的输出
boxes, classes, scores = prediction

# 将 boxes 和 scores 转化为 numpy 数组
boxes = boxes.numpy()
scores = scores.numpy()

# 筛选出置信度大于 50% 的检测结果
mask = scores > 0.5
boxes = boxes[mask]
scores = scores[mask]

# 绘制检测结果
for box, score in zip(boxes, scores):
    x1, y1, x2, y2 = box
    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.putText(image, f'{score:.2f}', (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

显示图片
cv2.imshow('image', image)
cv2.waitKey(0)
cv2.destroyAllWindows()

这样就可以使用 YOLOv5 模型对测试图片进行预测并可视化检测结果。