场景图生成网络——RelTR（TPAMI2023）

一、ReITR概述

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

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

RelTR 的主要贡献如下：

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

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

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

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

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

二、ReITR安装使用

1. 安装

  下载 RelTR 仓库：

git clone https://github.com/yrcong/RelTR.git
cd RelTR

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

# 创建一个 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。

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

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

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

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

  目录结构如下：

RelTR
| 
│
└───data
│   └───vg
│       │   rel.json
│       │   test.json
│       │   train.json
|       │   val.json
|       │   images
│   └───oi
│       │   rel.json
│       │   test.json
│       │   train.json
|       │   val.json
|       │   images
└───datasets    
... 

2. 使用方法

推理

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

ckpt/checkpoint0149.pth

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

python inference.py --img_path $IMAGE_PATH --resume $MODEL_PATH

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

训练

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

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 -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 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 main.py --dataset oi --img_folder data/oi/images/ --ann_path data/oi/ --eval --batch_size 1 --resume ckpt/checkpoint0149_oi.pth