ViRED: Prediction of Visual Relations in Engineering Drawings

The vision encoder in the model processes an image $I_{doc}$ of size $3\times H_{\textrm{Image}}\times W_{\textrm{Image}}$ and produces an image embedding in either feature vector format $F_{\textrm{vision}}$ or feature map format with dimensions $C_{v}\times H_{v}\times W_{v}$, where $C$, $H$, and $W$ represent the channels, width, and height of the image and feature map. In cases where feature maps are produced, they are flattened and converted into feature vectors. To prevent bias, a pre-trained vision encoder called the Masked Autoencoder (MAE) [30] is utilized in this study. The Document Image Transformer (DiT) [31] is employed as our vision encoder. The vision encoder is only used once for each image, regardless of the number of objects or relationships in the image.

The object encoder efficiently maps a bounding box and its typing information into a vectorial embedding. We use a convolutional neural network (CNN) to encode the bounding boxes. The object’s bounding box is shown as a one-channel image $I_{obj}$ with the same dimensions as the engineering drawing image. Pixels inside the bounding box are shown as 1, while those outside are shown as 0. The bounding box image is then encoded using a three-layer CNN.

To inform the relation decoder about whether the embedding token represents a circuit or a table, we aggregate the vector embeddings of bounding boxes $F_{\textrm{mask}}$ with two learned embeddings $F_{\textrm{type}}$.

The relation decoder takes the encoded masks and vectorized images as input and predicts if there is a relation between each circuit-table pair. In detail, the relation decoder consists of two components: the fusion model and the relation prediction model.

The fusion model is employed for feature fusion between object masks and image features. Inspired by DETR [3], a transformer-decoder-based model is adopted. We modify the standard transformer decoder by eliminating the relative position embedding and causal mask because of the lack of order between objects. This transforms the decoder into a bidirectional transformer decoder. As shown in Fig. 2, each decoder layer consists of four parts. 1. The tokens are initially processed by a self-attention model as shown in Eq. 5, which enables the mask tokens to interact with each other. This step enables the objects to determine their relative positions. 2. Next, a cross-attention model is introduced, where the image embeddings are used as the key and value vector, and mask tokens are used as the query vector like Eq. 6. By utilizing the image-to-object cross-attention model, the mask tokens are modified through a combination of bounding box features and vision features. 3. Then, a feedforward layer updates all the mask tokens, and a dropout layer is used to improve the model’s generalizability. 4. Finally, a residual connection is added to every attention layer and feed-forward layer, in accordance with the typical transformer architecture [10].

After receiving the object tokens from the fusion model of the relation decoder, it is essential to identify the relationships between them using the relation prediction model. The combined object tokens are concatenated to generate $(N_{c}+N_{t})^{2}$ combination tokens. This process is depicted in Fig. 3. To streamline the training and inference procedures, redundant and unfeasible combinations are eliminated. The relation prediction model consists of a three-layer perceptron with ReLU activation and a linear projection layer. This model is utilized to convert hidden features into a two-dimensional logit output. Afterwards, the main objective of the relation prediction model is to determine whether a relationship exists between two objects.

In the pretraining phase, our method is trained and tested on the PubLayNet datasets with the same train and test splits as the original dataset. For the finetuning phase, we used a self-annotated dataset of engineering design diagrams containing 283 images, 4,566 entities, and 2,112 relations between entities. The datasets are randomly divided, with 90% assigned for training and 10% for evaluation.

The dimension $D$ of latent representation throughout the model pipeline is set to 768. For the vision encoder backbone, we adopt the DINOv2-B model [35]. The model consists of 12 layers of Transformer encoders, each containing a 12-head multi-head attention block. The input image is cropped into patches of size 14, which are then processed and encoded to a 768-dimensional feature vector. We finetuned the backbone network using the same MAE training method as DiT [31] based on its state-of-the-art performance on visual-based document downstream tasks. The bounding box encoder is a three-layer convolution network with a ReLU activation function and a linear projection layer that maps the feature map to a 768-dimensional vectorized feature. The relation decoder is a 2-layer Transformer Decoder for multi-modal feature fusion. The input image size for the vision encoder is $518\times 518$, while the object encoder takes $224\times 224$ masks as input. Despite the difference in input sizes, we maintain the relative positions between the images and the masks to help the model better understand the positional relations between them.

We implement our model in PyTorch. For the pretraining phase, our model is trained with a batch size of 64 for a total of 100 epochs on the PubLayNet dataset using $\sim$90 NVIDIA A800 GPU hours. During the finetuning phase, we initialized the type embeddings to zero and the multi-layer perceptron of the relation extractor with uniformly random numbers. Meanwhile, other parameters remained unchanged from the pretraining checkpoint.

We applied data augmentation to the data while training, including horizontal and vertical flipping with a probability of 0.2, $D_{4}$ dihedral group transformations, and fixed-size random cropping. Unlike common random cropping, we ensured that instances in the images were not lost due to cropping.

For optimization, we use AdamW with hyper-parameters $\beta=(0.9,0.999),\epsilon=10^{-8}$ and weight decay of $10^{-3}$. The base learning rate for the optimizer is $10^{-4}$, we employ a linear learning rate schedule from $10^{-4}$ to $10^{-5}$ with 20% epoches for warming up.

Qualitative results are shown in Fig. 4. The blue and green areas represent the bounding box regions of circuits and tables, respectively. In the presence of a relationship between a table and a circuit, a line segment will link the two entities.

In this section, we evaluate our model against current methods for relation detection. We focus on visual relation detection instead of text-based methods for document relation extraction. To align with the visual task, we modify the dataset by categorizing textual descriptions as either ”table” or ”circuit.” Relations are established between all circuit-table pairs, with existing relations labeled as ”describe” and non-existing relations labeled as ”not describe.”

Table I provides a comparative analysis of our relation prediction method in contrast to several established approaches, with the accuracy metric used to measure the precision of the outcomes. Our findings demonstrate that our method surpasses the results of our engineering drawing relation prediction dataset. Importantly, we employ the evaluation methods used by the compared approaches to maintain consistency with our metrics. For MBBR, the accuracy is calculated based on the top-100 relation predictions.

In this section, we evaluate our method in different settings and strategies.

We conduct an additive ablation study on the mAP, precision, and recall for relation prediction on the validation set, as illustrated in Tab. III. In the baseline model, we initialize the model with random parameters without pretraining and omit the type embedding from the object encoder. Introducing type embedding allows the model to better differentiate the categories of object tokens. When computing self-attention over object tokens, the relation decoder can simultaneously consider the categories of the objects, thereby reducing computations for improbable relations and enhancing the understanding of relations. Due to the insufficient data in our task’s dataset, we introduced a pretraining mechanism for both the vision encoder and the position encoder-decoder. By undergoing unsupervised training on a large-scale document image dataset, the vision encoder more effectively extracts visual features from electrical engineering drawings. Supervised pretraining was utilized to predict the classification of selected regions based on given positional information, thereby enhancing the model’s understanding of positional inputs. As shown in Table III, the inclusion of type embedding and model pretraining significantly improved the performance across all evaluation metrics.

We also conducted ablation experiments on the model architecture. We compared the effects of various vision encoders, diverse object encoders, different input image resolutions, and different numbers of transformer decoder layers in the relation decoder on the performance. The experimental results are shown in Tab. II. The hidden dimension of the model is determined by the output dimension of the vision encoder. For ViT-S, ViT-B, ViT-L [35], and ResNet-50 [36], the dimensions for latent representation are set to 384, 768, 1024, and 768, respectively. Upon analyzing the experimental results, it becomes evident that the object encoder architecture significantly influences model performance. The relation decoder struggles to interpret object tokens encoded by ViT, resulting in a notable decline in performance across various metrics. Additionally, the vision encoder impacts the latent representation dimension of the model, thereby also playing a crucial role in determining model performance. The findings suggest that ViT-S performs poorly due to its constrained embedding dimension. For ViT-B and ViT-L, the limited task difficulty results in negligible performance differences. Higher image resolutions improve the vision encoder’s capacity to capture more detailed image features. However, the vision encoder’s ability to encode all these details is restricted by the dimension of the image feature vector, leading to a minimal effect on performance.

To evaluate the inference efficiency of the model, we compute the FLOPs (Floating Point Operations) for the inference process. For uniformity, a batch size of 1 is employed for all models. The input image resolution is fixed at $518\times 518$, with the number of objects to be predicted ($N$) ranging from 1 to 20. The relations to be predicted include all possible interactions between objects, amounting to $(N-1)^{2}$. The results are presented in Fig. 5. Owing to the utilization of a lightweight object encoder and relation decoder, the inference efficiency of ViRED is minimally impacted by the number of objects in the electrical engineering drawing. It sustains a rapid inference speed even when a single drawing contains numerous objects.