本人用yolov5训练出了一个猪识别模型，有什么代码可以调用权重文件best.pt进一步实现猪的身份识别？或者有什么类似的实操案例？

Question

本人用yolov5训练出了一个猪识别模型，有什么代码可以调用权重文件best.pt进一步实现猪的身份识别？或者有什么类似的实操案例？

Answer 1

基于Monster 组和GPT的调写：

用以下代码：

import torch
import cv2
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.datasets import letterbox

# 加载模型
model = attempt_load('weight.pt', map_location=torch.device('cpu'))

# 设置模型参数
img_size = 640
conf_thres = 0.25
iou_thres = 0.45

# 打开图像
img_path = 'test.jpg'
img = cv2.imread(img_path)

# 调整图像尺寸
img0 = img.copy()
img = letterbox(img, new_shape=img_size)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)

# 转换为Tensor
img = torch.from_numpy(img).to(device='cpu')
img = img.float()
img /= 255.0

# 添加batch size维度
img = img.unsqueeze(0)

# 进行推理
pred = model(img)

# 进行非极大值抑制
pred = non_max_suppression(pred, conf_thres, iou_thres)[0]

# 对结果进行后处理
if pred is not None:
    # 将预测框的坐标转换为原始图像的坐标
    pred[:, :4] = scale_coords(img.shape[2:], pred[:, :4], img0.shape).round()

    # 将预测框画在图像上
    for *xyxy, conf, cls in pred:
        label = f'{model.names[int(cls)]} {conf:.2f}'
        img0 = cv2.rectangle(img0, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (255, 0, 0), 2)
        img0 = cv2.putText(img0, label, (int(xyxy[0]), int(xyxy[1] - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
    
    # 显示图像
    cv2.imshow('result', img0)
    cv2.waitKey(0)

---

• 用了yolov5的Python实现，其中：
• 模型的权重文件为weight.pt
• 输入图像的尺寸为640x640
• 置信度阈值为0.25，IoU阈值为0.45

• 根据自己的实际情况调整参数。