实时开放词汇目标检测(论文复现)
本文所涉及所有资源均在传知代码平台可获取
文章目录
- 实时开放词汇目标检测(论文复现)
- 概述
- 模型框架
- 使用方式
- 配置环境
- 训练和评估
- 训练
- 评估
- 演示效果
- Gradio Demo
概述
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模型
文章代码资源点击附件获取