场景图生成网络——RelTR(TPAMI2023)

一、ReITR概述

场景图是一种图结构,其节点代表图像中的实体,边代表实体间的关系。这项技术超越了传统的对象检测,与视觉关系检测紧密相关,对图像检索、图像字幕、视觉问答(VQA)和图像生成等多种视觉-语言任务具有潜在的应用价值。

不同于传统的两阶段方法(先检测对象再分类关系),RelTR 采用一种一阶段方法,直接基于视觉外观预测场景图,无需合并实体或标记所有可能的谓词。这种方法基于DETR(Detection Transformer)的启发,将对象检测视为一种端到端的集合预测任务,并提出了一种基于二分匹配的集合预测损失函数。

RelTR 的主要贡献如下:

  • 相比大多数现有方法使用的对象检测骨架分类密集关系,RelTR 能够通过解码视觉外观及从数据中学习的主体和客体查询来生成稀疏场景图。

  • RelTR 仅基于视觉外观生成场景图,具有较少的参数和更快的推理速度,同时达到了最先进的性能。

  • 设计了一种集合预测损失来执行真实三元组和预测三元组之间的匹配,采用基于IoU的分配策略。

  • 通过解耦实体注意力,RelTR 的三元组解码器能够改进主体和客体的定位和分类。

论文还对在 Visual Genome、Open Images V6 和 VRD 数据集上的实验结果进行了详细讨论,证明了模型的优越性能和快速推理能力。RelTR 的源代码和预训练模型已公开提供,以便研究人员和开发者使用和进一步开发。

二、ReITR安装使用

  1. 安装

下载 RelTR 仓库:

python 复制代码
git clone https://github.com/yrcong/RelTR.git
cd RelTR

如果您想对一张图片进行推断,只需要 python=3.6, PyTorch=1.6 和 matplotlib!您可以按照以下方式配置环境:

python 复制代码
# 创建一个 conda 环境
conda create -n reltr python=3.6
conda activate reltr

# 安装包
conda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1 -c pytorch
conda install matplotlib

如果您想在 Visual Genome 上训练/评估 RelTR,您需要更多的准备:

a) 需要 Scipy(我们使用的是 1.5.2)和 pycocotools。

python 复制代码
conda install scipy
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

b) 按照 data 目录中的 README 准备数据集。

c) 一些常用的评估代码(IoU)需要编译... 我们将用 Pytorch 代码替换它。

python 复制代码
# 编译计算盒交集的代码
cd lib/fpn
sh make.sh

目录结构如下:

python 复制代码
RelTR
| 
│
└───data
│   └───vg
│       │   rel.json
│       │   test.json
│       │   train.json
|       │   val.json
|       │   images
│   └───oi
│       │   rel.json
│       │   test.json
│       │   train.json
|       │   val.json
|       │   images
└───datasets    
... 
  1. 使用方法

推理

a) 下载我们在 Visual Genome 数据集上预训练的 RelTR 模型,并将其放在

python 复制代码
ckpt/checkpoint0149.pth

b) 使用以下命令推断图片中的关系:

python 复制代码
python inference.py --img_path $IMAGE_PATH --resume $MODEL_PATH

我们附带了 VG 数据集的 5 张图片和 1 张互联网图片。您也可以用您自定义的图片进行测试。结果应该如下所示:

训练

a) 在单节点上使用 8 个 GPU 在 Visual Genome 上训练 RelTR(每个 GPU 2 张图片):

python 复制代码
python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --dataset vg --img_folder data/vg/images/ --ann_path data/vg/ --batch_size 2 --output_dir ckpt

b) 在单节点上使用 8 个 GPU 在 Open Images V6 上训练 RelTR(每个 GPU 2 张图片):

python 复制代码
python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --dataset oi --img_folder data/oi/images/ --ann_path data/oi/ --batch_size 2 --output_dir ckpt

评估

a) 使用单个 GPU 评估 Visual Genome 上的预训练 RelTR(每个 GPU 1 张图片):

python 复制代码
python main.py --dataset vg --img_folder data/vg/images/ --ann_path data/vg/ --eval --batch_size 1 --resume ckpt/checkpoint0149.pth

b) 使用单个 GPU 评估 Open Images V6 上的预训练 RelTR(每个 GPU 1 张图片):

python 复制代码
python main.py --dataset oi --img_folder data/oi/images/ --ann_path data/oi/ --eval --batch_size 1 --resume ckpt/checkpoint0149_oi.pth
相关推荐
说私域8 分钟前
基于开源AI智能名片链动2+1模式的S2B2C商城小程序:门店私域流量与视频号直播融合的生态创新研究
人工智能·小程序·开源
Ronin-Lotus11 分钟前
深度学习篇---Yolov系列
人工智能·深度学习
爱学习的茄子12 分钟前
AI驱动的单词学习应用:从图片识别到语音合成的完整实现
前端·深度学习·react.js
静心问道39 分钟前
GoT:超越思维链:语言模型中的有效思维图推理
人工智能·计算机视觉·语言模型
aneasystone本尊1 小时前
学习 Claude Code 的工具使用(三)
人工智能
szxinmai主板定制专家1 小时前
【精密测量】基于ARM+FPGA的多路光栅信号采集方案
服务器·arm开发·人工智能·嵌入式硬件·fpga开发
T__TIII1 小时前
Dify 自定义插件
人工智能·github
快起来别睡了1 小时前
LangChain 介绍及使用指南:从“会聊天”到“能干活”的 AI 应用开发工具
人工智能
AI数据皮皮侠2 小时前
中国区域10m空间分辨率楼高数据集(全国/分省/分市/免费数据)
大数据·人工智能·机器学习·分类·业界资讯
静心问道2 小时前
大语言模型能够理解并可以通过情绪刺激进行增强
人工智能·语言模型·大模型