SAHI усиливает эффективность YOLOv5 по небольшим целям

глубокое обучение

Мало знаний, большой вызов! Эта статья участвует в "Необходимые знания для программистов«Творческая деятельность.

окрестности

  • ubuntu 18.04 64bit
  • sahi 0.8.4
  • yolov5 5.0
  • pytorch 1.7.1+cu101

предисловие

Обнаружение объектов и сегментация экземпляров на сегодняшний день являются наиболее важными областями применения в компьютерном зрении, и различные сети обнаружения объектов возникают бесконечным потоком.Однако остается много проблем в практическом применении обнаружения мелких объектов и рассуждений на изображениях большого размера.SAHI (Slicing Aided Hyper Inference) разработан, чтобы помочь разработчикам решать эти реальные проблемы, это облегченная библиотека машинного зрения, которую можно использовать без переобучения исходной модели обнаружения (в настоящее время поддерживаетyolov5иmmdetection) для повышения скорости обнаружения малых целей, иgpuИспользование ресурсов существенно не улучшилось.

установить сахи

Установка библиотеки очень проста и может использоваться напрямую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реализация класса

Связанное чтение

использованная литература

Категории