Yolov8实现瓶盖正反面检测

news2024/9/20 16:39:10

一、模型介绍

模型基于 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.cls
    boxes_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()