环境
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
可以明显看到图片上方的一些汽车目标没有被检测出来
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
可以看到,在同一张测试图片上,使用同样的模型,sahi
库可以让 yolov5
检测出更多的目标
sahi
的切片推理原理如下图
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
类的实现
关联阅读
YOLOv5的5.0版本
pip安装YOLOv5
最后,国庆7天长假马上就要到了,提前祝小伙伴们节日快乐!