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,可能会打不开,后续会更新放上来,敬请期待!!!