本人用yolov5训练出了一个猪识别模型,有什么代码可以调用权重文件best.pt进一步实现猪的身份识别?或者有什么类似的实操案例?
用以下代码:
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)
训练代码(训练中。。。)
完成之后会在最新的exp下有个 weights \ bast.pt 文件,这个就是我们训练出来的模型
(CPU训练大概需要十七个小时左右,作者是使用GPU训练)
(1060显卡训练大概需要两小时,时间关系就没有完全训练)
引用chatGPT作答,恭喜您训练出了一个猪识别模型!要使用训练好的权重文件进行猪的身份识别,您可以使用以下代码片段:
import torch
import cv2
import numpy as np
# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='best.pt')
# 加载图片
img = cv2.imread('test.jpg')
# 进行推理
results = model(img)
# 获取结果
labels = results.xyxy[0][:, -1].numpy()
confidences = results.xyxy[0][:, 4].numpy()
boxes = results.xyxy[0][:, :4].numpy()
# 打印结果
for label, confidence, box in zip(labels, confidences, boxes):
print(f"Label: {label}, Confidence: {confidence}, Box: {box}")
请注意,上述代码需要安装 torch、opencv-python 和 numpy 库。您还需要将 test.jpg 替换为您想要进行预测的图像路径。
上述代码假设您的 YOLOv5 模型使用了相同的类别标签和相似的训练数据,因此您需要相应地调整代码以适应您的模型。
类似的实操案例包括目标检测、图像分类、人脸识别等,都可以使用类似的代码片段进行实现。