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

文章目录

• 环境
• 前言
• 安装sahi
• yolov5检测
• sahi
• 添加新的检测模型

环境

• 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 

首先安装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'

执行上面代码，得到

可以看到，在同一张测试图片上，使用同样的模型，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 只支持 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 类的实现。

QQ767172261