重磅干货,第一时间送达
在本文中,你将学习如何使用切片辅助超推理 (SAHI) 检测数据集中的小物体。我们将介绍以下内容:
-
为什么很难检测小物体
-
SAHI 的工作原理
-
如何将 SAHI 应用于你的数据集,以及
-
如何评估这些预测的质量
它们很小
首先,检测小物体很难,因为小物体很小。物体越小,检测模型需要处理的信息就越少。如果汽车在远处,它可能只占据我们图像中的几个像素。就像人类难以辨别远处的物体一样,我们的模型在没有车轮和车牌等视觉可辨别特征的情况下更难识别汽车!
训练数据
模型的好坏取决于训练它们的数据。大多数标准物体检测数据集和基准都集中在中到大型物体上,这意味着大多数现成的物体检测模型并未针对小物体检测进行优化。
固定输入大小
物体检测模型通常采用固定大小的输入。例如,YOLOv8 在最大边长为 640 像素的图像上进行训练。这意味着当我们输入一张 1920x1080 大小的图像时,模型会在进行预测之前将图像下采样到 640x360,降低分辨率并丢弃小物体的重要信息。
理论上,你可以在较大的图像上训练模型以改进对小物体的检测。但实际上,这需要更多的内存、更多的计算能力和更耗费人力的数据集。
另一种方法是利用现有的物体检测,将模型应用于图像中固定大小的块或切片,然后将结果拼接在一起。这就是切片辅助超推理(Slicing-Aided Hyper Inference)背后的想法!
SAHI 的工作原理是将图像分成完全覆盖它的切片,并使用指定的检测模型对每个切片进行推理。然后将所有这些切片的预测合并在一起,以生成整个图像的一个检测列表。SAHI 中的“超”来自这样一个事实,即 SAHI 的输出不是模型推理的结果,而是涉及多个模型推理的计算的结果。
💡SAHI 切片可以重叠(如上图 GIF 所示),这有助于确保至少一个切片中有足够的物体可供检测。
使用 SAHI 的主要优势在于它与模型无关。SAHI 可以利用当今的 SOTA 对象检测模型以及明天的 SOTA 模型!
当然,天下没有免费的午餐。作为“超推理”的交换,除了将结果拼接在一起所需的处理之外,你还要运行检测模型的多次前向传递。
为了说明如何应用 SAHI 来检测小物体,我们将使用天津大学机器学习与数据挖掘实验室的 AISKYEYE 团队的 VisDrone 检测数据集。该数据集包含 8629 张图像,边长从 360 像素到 2000 像素不等。Ultralytics 的 YOLOv8l 将作为我们的基础物体检测模型。
我们将使用以下库:
如果你还没有安装这些库的最新版本,请安装。你需要 fiftyone>=0.23.8 才能从 Hugging Face Hub 加载 VisDrone:
pip install -U fiftyone sahi ultralytics huggingface_hub
然后只需要导入相关的模块:
import fiftyone as fo
import fiftyone.zoo as foz
import fiftyone.utils.huggingface as fouh
from fiftyone import ViewField as F
就这样,我们就可以加载数据了!我们将使用 FiftyOne 的 Hugging Face 实用程序中的 load_from_hub() 函数,通过其 repo_id 直接从 Hugging Face Hub 加载 VisDrone 数据集的一部分。
为了演示并尽可能快地执行代码,我们将仅从数据集中获取前 100 张图像。我们还将为这个正在创建的新数据集命名为“sahi-test”:
dataset = fouh.load_from_hub(
"Voxel51/VisDrone2019-DET",
name="sahi-test",
max_samples=100
)
在添加任何预测之前,让我们先看看 FiftyOne 应用程序中的数据集:
session = fo.launch_app(dataset)
在下一节中,我们将使用 SAHI 对我们的数据进行超推理。在引入 SAHI 之前,让我们使用 Ultralytics 的 YOLOv8 模型的大型变体对我们的数据进行标准对象检测推理。
首先,我们创建一个 ultralytics.YOLO 模型实例,并在必要时下载模型检查点。然后,我们将此模型应用于我们的数据集,并将结果存储在样本的字段“base_model”中:
from ultralytics import YOLO
ckpt_path = "yolov8l.pt"
model = YOLO(ckpt_path)
dataset.apply_model(model, label_field="base_model")
session.view = dataset.view()
通过查看模型的预测和真实值标签,我们可以看到一些事情。首先,我们的 YOLOv8l 模型检测到的类别与 VisDrone 数据集中的真实值类别不同。我们的 YOLO 模型是在 COCO 数据集上训练的,该数据集有 80 个类别,而 VisDrone 数据集有 12 个类别,包括一个 ignore_regions 类。
为了简化比较,我们将只关注数据集中最常见的几个类别,并将 VisDrone 类映射到 COCO 类,如下所示:
mapping = {"pedestrians": "person", "people": "person", "van": "car"}
mapped_view = dataset.map_labels("ground_truth", mapping)
然后过滤标签,仅包含我们感兴趣的类别:
def get_label_fields(sample_collection):
"""Get the (detection) label fields of a Dataset or DatasetView."""
label_fields = list(
sample_collection.get_field_schema(embedded_doc_type=fo.Detections).keys()
)
return label_fields
def filter_all_labels(sample_collection):
label_fields = get_label_fields(sample_collection)
filtered_view = sample_collection
for lf in label_fields:
filtered_view = filtered_view.filter_labels(
lf, F("label").is_in(["person", "car", "truck"]), only_matches=False
)
return filtered_view
filtered_view = filter_all_labels(mapped_view)
session.view = filtered_view.view()
现在我们有了基本模型预测,让我们使用 SAHI 来对图像进行切片和切块💪。
SAHI 技术在我们之前安装的 sahi Python 包中实现。SAHI 是一个与许多对象检测模型兼容的框架,包括 YOLOv8。我们可以选择想要使用的检测模型,并创建任何 sahi.models.DetectionModel 子类的实例,包括 YOLOv8、YOLOv5 甚至 Hugging Face Transformers 模型。
我们将使用 SAHI 的 AutoDetectionModel 类创建模型对象,并指定模型类型和检查点文件的路径:
from sahi import AutoDetectionModel
from sahi.predict import get_prediction, get_sliced_prediction
detection_model = AutoDetectionModel.from_pretrained(
model_type='yolov8',
model_path=ckpt_path,
confidence_threshold=0.25, ## same as the default value for our base model
image_size=640,
device="cpu", # or 'cuda' if you have access to GPU
)
在生成切片预测之前,让我们使用 SAHI 的 get_prediction() 函数检查模型在试验图像上的预测:
result = get_prediction(dataset.first().filepath, detection_model)
print(result)
输出是:
<sahi.prediction.PredictionResult object at 0x2b0e9c250>
幸运的是,SAHI 结果对象有一个 to_fiftyone_detections() 方法,它将结果转换为 FiftyOne Detection 对象列表:
print(result.to_fiftyone_detections())
这样我们就可以专注于数据,而不是繁琐的格式转换细节。SAHI 的 get_sliced_prediction() 函数的工作方式与 get_prediction() 相同,但增加了一些超参数,让我们可以配置图像的切片方式。具体来说,我们可以指定切片高度和宽度,以及切片之间的重叠。以下是示例:
sliced_result = get_sliced_prediction(
dataset.skip(40).first().filepath,
detection_model,
slice_height = 320,
slice_width = 320,
overlap_height_ratio = 0.2,
overlap_width_ratio = 0.2,
)
作为初步检查,我们可以将切片预测中的检测数量与原始预测中的检测数量进行比较:
num_sliced_dets = len(sliced_result.to_fiftyone_detections())
num_orig_dets = len(result.to_fiftyone_detections())
print(f"Detections predicted without slicing: {num_orig_dets}")
print(f"Detections predicted with slicing: {num_sliced_dets}")
Detections predicted without slicing: 17
Detections predicted with slicing: 73
我们可以看到预测数量大幅增加!我们尚未确定这些额外的预测是否有效,或者我们是否只是有更多的误报。我们将很快使用 FiftyOne 的评估 API 来做到这一点。我们还想为我们的切片找到一组好的超参数。我们需要将 SAHI 应用于整个数据集来完成所有这些事情。现在就开始吧!
为了简化流程,我们将定义一个函数,将预测添加到指定标签字段中的样本,然后我们将遍历数据集,将该函数应用于每个样本。此函数将样本的文件路径和切片超参数传递给 get_sliced_prediction(),然后将预测添加到指定标签字段中的样本:
def predict_with_slicing(sample, label_field, **kwargs):
result = get_sliced_prediction(
sample.filepath, detection_model, verbose=0, **kwargs
)
sample[label_field] = fo.Detections(detections=result.to_fiftyone_detections())
我们将切片重叠固定为 0.2,并观察切片高度和宽度如何影响预测的质量:
kwargs = {"overlap_height_ratio": 0.2, "overlap_width_ratio": 0.2}
for