SAHI强化YOLOv5在小目标上的表现

点击下方“AI算法与图像处理”，一起进步！
重磅干货，第一时间送达

环境

• 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 类的实现