实时开放词汇目标检测（论文复现）

本文所涉及所有资源均在传知代码平台可获取

概述

YOLO-World是由腾讯人工智能实验室于2024年1月31日发布的实时开放词汇目标检测模型，能够在实时环境中跨越开放词汇表识别对象，无需先前的训练。传统的目标检测模型如YOLO由于依赖于预定义和训练过的目标类别（闭集检测），它在开放场景中的适用性受到了限制，例如，使用COCO数据集训练的模型仅能识别80个不同的类别。为了应对固定词汇检测器的限制，开放词汇目标检测（OVD）的概念应运而生，旨在识别超出预先建立类别范围之外的对象。

YOLO-World利用大量的图像-文本对和基础图像进行训练，以理解和响应各种提示，例如“穿着黑色裤子的人”。通过引入“提示-然后检测”的方法论，YOLO-World避开了即时文本编码的需要，而是利用用户提示生成的离线词汇来进行检测。这种方法显著降低了计算需求，允许灵活调整检测词汇，以满足各种需求，而不会影响性能，从而拓展了模型在实际场景中的适用性。在LVIS这个具有挑战性的数据集上，YOLO-World在V100上达到了35.4的AP和52的FPS，无论是精度和速度上都超越了以前SOTA的方法，如下图所示

模型框架

上图所示为YOLO-World的整体框架，主要包括了YOLO检测器，Text Encoder，和RepVL-PAN（Re-parameterizable Vision-Language Path Aggregation NetWork）。与传统检测器不同的是，YOLO-World作为开集检测器，需要使用文本作为输入，Text Encoder首先会编码输入的文本，然后输出Vocabulary embedding；之后Image Encoder（backbone）会编码输入图像，或者说提取图像特征，以获得多尺度特征图；RepVL-PAN（Vision-Language PAN）会利用图像和文本特征的多层次跨模态进行融合；最后，YOLO-World会预测出回归框和目标embedding，去匹配在输入文本中的出现的类别或者名词。

YOLO检测器

YOLO-World 是基于YOLOv8开发出来的，它包含了Darknet的backbone作为图像encoder，一个路径聚合网络（PAN）构建多尺度特征金字塔，以及一个输出回归边界框和目标embedding的预测头。

文本encoder

给定文本T，我们使用预训练CLIP的Transformer text encoder抽取相关的文本embedding。CLIP的text encoder能够提供更好的视觉-语义能力，使得视觉目标和文本相互连接。

文本对比头(Text Contrastive Head)

使用了yolov8的解耦头和俩个3×3卷积。因为要计算目标-文本的相似度，所以提出文本对比头。为了稳定区域-文本训练，目标编码e和文本编码t使用L2-Norm。

在线词汇表

在训练过程中，为每个包含4幅图像的马赛克样本构建一个在线词汇表 。

离线词汇表

提出了一种以“提示后检测”的策略，使用离线词汇以进一步提高效率。离线词汇表，特指的是经过encoder的embedding，也就是类别名，名词短语和目标描述构成的特征矩阵。与之对应的是，在线词汇则表示的不是embedding，在线词汇指的是没有经过encoder编码后的词汇

RepVL-PAN

RepVL-PAN的内部结构如上图所示。其中，文本引导的CSPLayer（T-CSPLayer），负责将语言信息注入图像特征中；而图像池化注意力Image Pooling Attention（I-Pooling Attention）则是负责强化具备图像意识的text embedding，以进一步增强图像特征与文本特征之间的交互，这可以提高开集能力的视觉语义表示。RepVL-PAN使用的和YOLOv8相似的特征融合结构，包括了自上而下和自下而上的路径，使用了多尺度的图像特征 {C3,C4,C5}搭建了特征金字塔 {P3,P4,P5}

使用方式

配置环境

创建python虚拟环境并激活虚拟环境

conda create -n yolov_world python=3.10
conda activate yolov_world

安装依赖包

# 安装pytorch
pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

# 安装其他依赖，这一步会比较久
pip install -e .
# 如果一直卡在Building wheel for mmcv就尝试下面的两个命令然后再重试
pip install -U openmim
mim install mmcv

训练和评估

训练

使用mmyolo默认的训练脚本，位于configs/pretrain目录下。注意：YOLO-World是在4个节点（每个节点配有8个GPU，总计32个GPU）上训练的。

# 给脚本增加可执行权限
chmod +x tools/dist_train.sh
# 训练的示例命令,使用8个GPU，开启AMP混合精度训练
./tools/dist_train.sh configs/pretrain/yolo_world_l_t2i_bn_2e-4_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py 8 --amp

评估

使用mmyolo默认的评估脚本，位于configs/finetune_coco目录下。主要在LVIS-minival数据集上评估预训练模型的性能

# # 给脚本增加可执行权限
chmod +x tools/dist_test.sh
# 评估的命令
./tools/dist_test.sh path/to/config path/to/weights 8

演示效果

python image_demo.py path/to/config path/to/weights image/path/directory 'person,dog,cat' --topk 100 --threshold 0.005 --output-dir demo_outputs

# 将path/to/config替换成配置文件的路径，path/to/weights换成模型权重路径，模型权重可以在https://huggingface.co/wondervictor/YOLO-World/tree/main中下载，注意配置文件和模型权重需要对应，image/path/directory换成自己要检测图像的路径，'person,dog,cat' 换成自己感兴趣的类别，运行成功后结果在demo_outputs下查看。
# python image_demo.py configs/pretrain/yolo_world_v2_m_vlpan_bn_2e-3_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py ./yolo_world_v2_m_obj365v1_goldg_pretrain-c6237d5b.pth data/images 'person,dog,cat' --topk 100 --threshold 0.005 --output-dir demo_outputs

YOLO-World 框架允许通过自定义提示动态指定类别，使用户能够根据自己的特定需求定制模型，而无需重新训练。通过设置自定义提示，用户可以引导模型关注感兴趣的对象，从而提高检测结果的相关性和准确性。如将上面的’person,dog,cat’ 换成自己感兴趣的类别。

（注：如果运行报错：Incorrect path_or_model_id: ‘…/pretrained_models/clip-vit-base-patch32-projection’.将configs目录下对应的配置文件（如上面的configs/pretrain/yolo_world_v2_m_vlpan_bn_2e-3_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py ）中第一个text_model_name注释掉，第二个text_model_name取消注释，从而自动从huggingface下载clip模型。）

Gradio Demo

在本地机器上运行web界面

pip install gradio==4.16.0 onnx onnxsim # 安装依赖
python demo.py path/to/config path/to/weights

# python demo.py configs/pretrain/yolo_world_v2_m_vlpan_bn_2e-3_100e_4x8gpus_obj365v1_goldg_train_lvis_minival.py ./yolo_world_v2_m_obj365v1_goldg_pretrain-c6237d5b.pth

运行成功然后在浏览器访问http://127.0.0.1:8080可以看到下面的界面，左上方输入检测的图片，下方的文本框输入想检测的类别，然后点击submit就能在右边生成检测结果。下面可以调节一些参数，如NMS Threshold等。还支持导出onnx模型

文章代码资源点击附件获取