场景图生成网络——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
相关推荐
橡晟3 小时前
深度学习入门:让神经网络变得“深不可测“⚡(二)
人工智能·python·深度学习·机器学习·计算机视觉
墨尘游子3 小时前
神经网络的层与块
人工智能·python·深度学习·机器学习
Leah01053 小时前
什么是神经网络,常用的神经网络,如何训练一个神经网络
人工智能·深度学习·神经网络·ai
Leah01053 小时前
机器学习、深度学习、神经网络之间的关系
深度学习·神经网络·机器学习·ai
PyAIExplorer3 小时前
图像亮度调整的简单实现
人工智能·计算机视觉
Striker_Eureka4 小时前
DiffDet4SAR——首次将扩散模型用于SAR图像目标检测,来自2024 GRSL(ESI高被引1%论文)
人工智能·目标检测
Rvelamen5 小时前
LLM-SECURITY-PROMPTS大模型提示词攻击测评基准
人工智能·python·安全
AI technophile5 小时前
OpenCV计算机视觉实战(15)——霍夫变换详解
人工智能·opencv·计算机视觉
JNU freshman6 小时前
计算机视觉 之 数字图像处理基础(一)
人工智能·计算机视觉
鹧鸪云光伏6 小时前
鹧鸪云重构光伏发电量预测的精度标准
人工智能·无人机·光伏·光伏设计·光伏模拟