环境
- ubuntu 18.04 64bit
- sahi 0.8.4
- yolov5 5.0
- pytorch 1.7.1+cu101
前言
目标检测和实例分割是迄今为止计算机视觉中最重要的应用领域,各种目标检测网络层出不穷,然而,小目标的检测和在大尺寸图像上的推理在实际应用中依然有诸多问题。 SAHI
(Slicing Aided Hyper Inference
)就是用来帮助开发人员解决这些现实问题,它是一个轻量级的视觉库,可以在不重新训练原始检测模型(目前支持 yolov5
和 mmdetection
)的情况下提升小目标的检出率,而且 gpu
资源的使用并没有大幅提升。
安装sahi
库的安装非常简单,直接使用 pip
安装最新版本,执行命令
pip install sahi
yolov5检测
sahi
使用的 YOLOv5, 这个我们前文已经介绍过了,参考 pip安装YOLOv5
# 安装gpu版torch和torchvision,这里使用1.7.1版本
pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
# 安装其它依赖
pip install yolov5
找到 sahi
项目中的一张图片来测试
# 下载模型
wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s6.pt
# 下载测试图片
wget https://raw.githubusercontent.com/obss/sahi/main/demo/demo_data/small-vehicles1.jpeg
# 命令行检测
yolov5 detect --source small-vehicles1.jpeg --weights yolov5s6.pt
可以明显看到图片上方的一些汽车目标没有被检测出来
sahi
首先安装2个基础依赖库
pip install fiftyone imantics
再来看下面的示例代码
from sahi.utils.yolov5 import download_yolov5s6_model
from sahi.model import Yolov5DetectionModel
from sahi.utils.file import download_from_url
from sahi.predict import get_sliced_prediction
yolov5_model_path = 'yolov5s6.pt'
# 下载模型
download_yolov5s6_model(destination_path=yolov5_model_path)
# 下载测试图片
download_from_url('https://raw.githubusercontent.com/obss/sahi/main/demo/demo_data/small-vehicles1.jpeg', 'small-vehicles1.jpeg')
# 使用的YOLOv5检测模型,使用gpu加速,置信度0.3
detection_model = Yolov5DetectionModel(
model_path=yolov5_model_path,
confidence_threshold=0.3,
device="cuda", # or 'cpu'
)
# 方法将待测试图片分成多个小图(默认是256x256),各自分别检测,最后进行拼接。小图框高宽重叠默认0.2,换算成像素就是256x0.2=51pixel。如果需要检测图片文件夹的话,可以使用方法predict
result = get_sliced_prediction(
"small-vehicles1.jpeg",
detection_model,
slice_height = 256,
slice_width = 256,
overlap_height_ratio = 0.2,
overlap_width_ratio = 0.2
)
# 保存检测图片
result.export_visuals(export_dir="result/")
执行上面代码,得到
可以看到,在同一张测试图片上,使用同样的模型,sahi
库可以让 yolov5
检测出更多的目标
sahi
的切片推理原理如下图
sahi
提供了命令行工具 sahi
,使用它可以完成快速检测
sahi predict --source small-vehicles1.jpeg --model_type yolov5 --model_path yolov5s6.pt --model_device cuda --export_visual --project results
更多参数及使用帮助,参考 sahi predict --help
添加新的检测模型
目前的 sahi
只支持 yolov5
和 mmdetection
,但是,我们可以很方便的添加新框架的支持,你只需要在 sahi/sahi/model.py
文件中创建一个新的类,这个类继承自 DetectionModel
,然后依次实现 load_model
、perform_inference
、_create_object_prediction_list_from_original_predictions
、_create_original_predictions_from_object_prediction_list
这几个方法即可,具体的可以参考 Yolov5DetectionModel
类的实现