一、模型介绍
- 模型基于 yolov8n
- 数据集采用SKU-110k,这数据集太大了十几个 G,所以只训练了 10 轮左右就拿来微调了
- 基于原木数据微调:训练 200 轮的效果
10 轮SKU-110k + 20 轮原木+ 200 轮瓶盖正反面 微调模型下载地址https://wwxd.lanzouu.com/iREN21gey34d
密码:d8x3下载完把模型格式修改为 pt
二、模型测试
- 废话不多说上代码:
import cv2
from ultralytics import YOLO# 加载模型
model = YOLO('best_bottle.pt') # 预训练的 YOLOv8n 模型image_url = '123.webp'
# 在图片列表上运行批量推理
results = model([image_url]) # 返回 Results 对象列表image = cv2.imread(image_url) # 替换成实际图像路径# 获取类别映射
class_mapping = {0: "BACK", 1: "FRONT"}# 处理结果列表
for result in results:boxes = result.boxes # 边界框输出的 Boxes 对象labels = boxes.clsboxes_np = boxes.xyxy.numpy()for box, label in zip(boxes_np, labels):x1, y1, x2, y2 = map(int, box)class_index = label.item()class_label = class_mapping.get(class_index, "Unknown")# 设置颜色color = (0, 255, 0) if class_label == "BACK" else (0, 0, 255)# 在图像上绘制边界框image = cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)# 在图像上绘制分类标签label_text = f"{class_label}"label_size, _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)label_position = (x1 + int((x2 - x1) / 2), y1 + int((y2 - y1) / 2))cv2.putText(image, label_text, label_position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)# 创建窗口并展示图像
cv2.namedWindow('Bounding Boxes', cv2.WINDOW_NORMAL)
cv2.imshow('Bounding Boxes', image)
cv2.waitKey(0)
cv2.destroyAllWindows()