【YOLO】YOLOv8 + SAHI 小目标检测（持续更新）

https://github.com/obss/sahi

SAHI小目标检测切片辅助

0.准备

• 下载最新 SAHI ultralytics:

pip install -U torch sahi ultralytics

import os
# os.getcwd()

• 导入模块

# arrange an instance segmentation model for test
from sahi.utils.yolov8 import (
    download_yolov8s_model, download_yolov8s_seg_model
)

from sahi import AutoDetectionModel
from sahi.utils.cv import read_image
from sahi.utils.file import download_from_url
from sahi.predict import get_prediction, get_sliced_prediction, predict
from IPython.display import Image

• 识别模型

# yolov8模型(可以是通过yolov8 自己训练好的模型)
yolov8_model_path = "models/yolov8_a6000_best.pt"

1.使用 YOLOv8 模型进行标准推理

• 通过定义模型权重、路径和其他参数来实例化检测模型：

detection_model = AutoDetectionModel.from_pretrained(
    model_type='yolov8',
    model_path=yolov8_model_path,
    confidence_threshold=0.4,
    device='cuda:0',  # or "cpu"
)

• 通过向 get_prediction 函数提供图像路径和 DetectionModel 实例来执行预测：

result = get_prediction("demo_data/test.jpg", detection_model)

• 或者，通过向 get_prediction 函数提供 numpy 图像和 DetectionModel 实例来执行预测：

result = get_prediction(read_image("demo_data/test.jpg"), detection_model)

• 在原始图像上可视化预测的边界框和掩码：

result.export_visuals(export_dir="demo_data/")

Image("demo_data/prediction_visual.png")

2.使用 YOLOv8 模型进行切片推理

• 为了执行切片预测，我们需要指定切片参数。在此示例中，我们将对重叠率为 0.05 的 1080x640 的切片进行预测：

result = get_sliced_prediction(
    "demo_data/test.jpg",
    detection_model,
    slice_height=1080,
    slice_width=640,
    overlap_height_ratio=0.05,
    overlap_width_ratio=0.05
)

• 在原始图像上可视化预测的边界框和掩码：

result.export_visuals(export_dir="demo_data/", hide_labels=True, hide_conf=True)

Image("demo_data/prediction_visual.png")

3. 批量预测

• 设置模型和资源目录（存图片的文件夹）参数：

model_type = "yolov8"
# 自己训练的模型地址
yolov8_model_path = "models/yolov8_a6000_best.pt"
model_device = 'cuda:0'  # or "cpu"
model_confidence_threshold = 0.45

slice_height = 1080
slice_width = 640
overlap_height_ratio = 0.05
overlap_width_ratio = 0.05

# 识别图片文件夹
source_image_dir = "data/2024_06/"

• 预测文件夹下所有图片

predict(
    model_type=model_type,
    model_path=yolov8_model_path,
    model_device=model_device,
    model_confidence_threshold=model_confidence_threshold,
    source=source_image_dir,
    slice_height=slice_height,
    slice_width=slice_width,
    overlap_height_ratio=overlap_height_ratio,
    overlap_width_ratio=overlap_width_ratio,
    visual_hide_conf=True, # 隐藏置信度值（默认False）
    visual_hide_labels=True, # 隐藏标签名（默认False）
    visual_bbox_thickness=2, # 框的描边粗细（int值）
)

结果导出

result.export_visuals(export_dir="demo_data/", hide_conf=True, hide_labels=True) 
# Image("demo_data/prediction_visual.png")

将yolov8模型转换格式

将pt转化为onnx

from ultralytics import YOLO


class Yolov8Pt2Onnx:
    def __init__(self, pt_path) -> None:
        self.convert(pt_path)

    def convert(self, pt_path):
        model = YOLO(pt_path)
        model.export(format="onnx", opset=12)


if __name__ == "__main__":
    pt_path = "./best.pt"
    Yolov8Pt2Onnx(pt_path)

.onnx模型可用于自动化数据标注x-anyLabeling

自动标注工具参考我的另一篇博客：https://blog.csdn.net/steamedobun/article/details/141640321

参数解析

1.source: 输入源，可以是图像文件路径、视频文件路径、摄像头索引等。
2.model: 预训练模型的路径或模型对象。
3.conf: 置信度阈值，滤除低于该置信度的检测结果。
4.iou: NMS（非极大值抑制）的交并比阈值，用于消除重叠的检测框。
5.device: 运行推理的设备，比如 ‘cpu’ 或 ‘cuda’。
6.imgsz: 输入图像的尺寸，可以是单个整数（长宽相等）或元组（长和宽）。
7.half: 是否使用半精度浮点数进行推理，通常用于加速推理过程。
8.save_txt: 是否将预测结果保存为文本文件。
9.save_conf: 是否在保存的文本文件中包含置信度。
10.save_crop: 是否保存裁剪后的检测结果。
11.save_img: 是否保存带有检测框的图像。
12.classes: 指定要检测的类的索引列表。
13.agnostic_nms: 是否使用类别无关的 NMS。
14.augment: 是否使用数据增强。
15.visualize: 是否可视化特征图。
16.project: 保存预测结果的项目目录。
17.name: 保存预测结果的子目录名称。
18.exist_ok: 是否允许项目目录存在。
19.line_thickness: 绘制检测框的线条粗细。
20.hide_labels: 是否隐藏检测框的标签。
21.hide_conf: 是否隐藏检测框的置信度。

SAHI 还支持以下模型辅助

官网：https://github.com/obss/sahi

这里用到了Google Colab，可能会打不开，后续会更新放上来，敬请期待！！！