Improving Transformer Based Line Segment Detection with Matched
Predicting and Re-ranking

In traditional DETR frameworks, bipartite matching is used to assign predictions to ground truth sequences. However, these approaches can be unstable in convergence, particularly during the early stages of training (Li et al. 2022). For the task of line segment detection, we observe that the centroids of individual line segments often do not coincide. Based on this observation, we introduce a novel centroid-based representation for line segments in Transformer-based predictions. Specifically, in proposal-based line segment detection, each feature point in the feature map has a unique 2D coordinate, and the feature point closest to the centroid of a line segment is responsible for predicting it. Leveraging the Transformer architecture, which incorporates positional information through positional encoding, each feature point can directly predict the endpoint coordinates $(x_{s},y_{s},x_{e},y_{e})$ of the line segment, rather than predicting the offsets of the endpoints from the centroid as in CNN-based methods (Huang et al. 2020b; Gu et al. 2022).

Given the processed feature point $\bm{p}$ in the feature map $\bm{F}$, we can get the confidence score $c$ indicating the confidence that the feature point should predict a line segment

$$c=\sigma(\mathcal{F}_{c}(\bm{p})),$$

(1)

and the location of the corresponding line segment

$$\bm{l}=\mathcal{F}_{l}(\bm{p}),$$

(2)

where $\mathcal{F}_{c},\mathcal{F}_{l}$ are CNN-based prediction modules and $\sigma$ is the sigmoid function.

While confidence scores are assigned to predicted line segments, these scores may not fully reflect the true quality of the segments since they are generated simultaneously. We propose optimizing these scores by evaluating the predicted line segments against corresponding geometric information, such as edges, endpoints, and length. To accomplish this, we perform sampling-based evaluations separately for endpoints and edges using the generated line segments. Specifically, we create endpoint maps and edge maps with a geometric information extractor. For a line segment to be considered high-quality, its endpoints should align with high-confidence regions on the endpoint map, and the segment should significantly overlap with high-confidence areas on the edge map.

Given the endpoint map $\bm{M_{e}}$, we sample the endpoints $\bm{e_{1}},\bm{e_{2}}$ according to the predicted line segments. The sampled scores are averaged as the endpoint confidence, which can be represented as

$$s_{e}=\frac{1}{2}\sum\mathcal{S}(\{\bm{e_{1}},\bm{e_{2}}\};\bm{M_{e}})),$$

(3)

$\mathcal{S}(\{\bm{e}\};\bm{M})$ represents the sampling kernel indicating sampling $\bm{M}$ with the location $\bm{e}$. Given the edge map $\bm{M_{d}}$, we first uniformly sample $m$ points between the two endpoints of the line segments. The scores are then sampled with these points from the edge map, and finally averaged as the edge score, which can be represented as

$$s_{d}=\frac{1}{m}\sum\mathcal{S}(\{\frac{k}{m}\bm{e_{1}}+\frac{m-k}{m}\bm{e_{2}}\}_{k\in[1,m]};\bm{M_{d}})).$$

(4)

The geometric information extractor is composed of a group of CNNs in this work, which is supervised by ground truth edge maps and endpoint maps generated from the ground truth line segments.

Moreover, to encourage the detection of long line segments and suppress the detection of fragmented line segments, the length is taken into account in evaluating line segment quality. To avoid the score expansion caused by length, the length scores is defined as

$$s_{l}=ln(\|\bm{e_{1}}-\bm{e_{2}}\|_{2}+1).$$

(5)

The optimized score is finally defined as the sum of the above terms and the previous confidence score $c$, which can be represented as

$$s=\delta_{e}s_{e}+\delta_{d}s_{d}+\delta_{l}s_{l}+c.$$

(6)

By incorporating the re-ranking method, we can refine the confidence scores of predicted line segments by considering not only the centroid confidence score but also the overall quality of the line segments, as indicated by their endpoints, overlap with the edge map and the length.

Our network architecture is built on a CNN-based backbone combined with a Deformable Transformer, designed specifically for line segment detection, which requires abundant low-level, high-resolution information. To capture features at various resolutions, we employ a feature pyramid during CNN feature extraction, feeding these multi-scale features into the Transformer. Notably, our approach simplifies the Transformer by using only the Encoder befitting from the stable line segment representation. Each feature token in the Encoder layers corresponds directly to a prediction target, eliminating the need for learnable queries and bipartite matching. For line segment prediction, we utilize the final $128\times 128$ high-resolution features, which allow for finer predictions and reduce conflicts from overlapping centroids. Additionally, we enhance the network’s sensitivity to line segments by incorporating rotation augmentation during image feature extraction, rotating the images by ±90° and applying an inverse rotation afterward. Our overall architecture includes an image feature extractor based on ResNet50 with a feature pyramid and a feature processing module consisting of 6 deformable Transformer Encoder layers.

Since our method maps predicted results to ground truth before prediction, bipartite matching is no longer required to establish correspondences. Instead, we directly employ confidence loss and position loss to supervise our line segment detection model during training. Specifically, for confidence loss, we use binary cross-entropy loss to determine whether a feature point should predict a line segment, i.e., whether the centroid of the line segment is closest to the feature point. We find that the number of feature points that do not need to be responsible for predicting the line segments (defined as negative feature points) is considerably larger than those that need to be responsible (defined as positive feature points). Thus, we randomly select some of the former for training. The confidence loss $\mathcal{L}_{conf}$ is defined as

$$\mathcal{L}_{conf}=-\frac{1}{N_{\pm}}\sum_{i=1}^{N_{\pm}}{(\hat{c_{i}}\log c_{i}+(1-\hat{c_{i}})\log(1-c_{i}))},$$

(7)

where $N_{\pm}$ is the total number of the positive and selected negative feature points. $c_{i}$ is the predicted confidence score of the $i$-th feature point. $\hat{c_{i}}$ is the ground truth indicates that whether the $i$-th feature point should respond to detect a line segment and is annotated as $+1$ or $0$.

In the position loss, we minimize the L1 distance between the predicted line segment and ground truth ones on the image space. The loss is only applied to the positive feature points since the prediction of negative anchors should be meaningless. The presented position loss $\mathcal{L}_{pos}$ can be represented as

$$\mathcal{L}_{pos}=\sum_{i=1}^{N_{+}}{\|\bm{l_{i}}-\bm{\hat{l_{i}}}\|_{1}},$$

(8)

where $N_{+}$ is the number of positive feature points. $\bm{l_{i}}$ represents a predicted line segments from the positive feature point and $\bm{\hat{l_{i}}}$ is the corresponding ground-truth.

Inspired by the success of recent ranking-based losses in detection tasks, we introduce a line segment ranking loss to ensure that feature points with higher line segment detection quality receive higher confidence scores. Many ranking-based losses are guided by the quality of the bounding box or the predicting confidence. In the context of line segment detection, the natural basis for ranking is the quality of the predicted line segment. We define this quality using the Euclidean distance between the predicted line segment and its corresponding ground truth. A shorter distance indicates higher prediction quality. For all positive feature points, we perform pairwise comparisons based on their prediction quality, ensuring that feature points with higher prediction quality receive higher confidence scores, which can be represented as

	$\displaystyle\mathcal{L}_{rank}=$	$\displaystyle-\frac{1}{N_{+}^{2}}\sum_{i=1}^{N_{+}}\sum_{j=1}^{N_{+}}\sigma((c_{j}-c_{i}))$		(9)
		$\displaystyle(\|\bm{l_{i}}-\bm{\hat{l_{i}}}\|_{2}-\|\bm{l_{j}}-\bm{\hat{l_{j}}}\|_{2}),$

where we use sigmoid function to approximate the step function.

Moreover, as we need to use the junction map and edge map in our line re-ranking module, we predict junction maps and edge maps from feature maps in different resolutions. For each resolution, we supervise the junction map $\mathcal{F}_{j}(\bm{F})$ and edge map $\mathcal{F}_{e}(\bm{F})$ as similar to ground truth maps $J$ and $E$, which can be represented as

$$\mathcal{L}_{junc}=\sum_{i=1}^{n}{\|\mathcal{F}_{j}(\bm{F_{i}})-J_{i}\|_{2}},$$

(10)

and

$$\mathcal{L}_{edge}=\sum_{i=1}^{n}{\|\mathcal{F}_{e}(\bm{F_{i}})-E_{i}\|_{2}}.$$

(11)

The total loss can be defined as the sum of the above loss terms with proper weights, which can be written as

	$\displaystyle\mathcal{L}_{total}=$	$\displaystyle\lambda_{r}\mathcal{L}_{rank}+\lambda_{c}\mathcal{L}_{conf}+\lambda_{p}\mathcal{L}_{pos}+\lambda_{j}\mathcal{L}_{junc}$		(12)
		$\displaystyle+\lambda_{e}\mathcal{L}_{edge}.$

Our comparisons are conducted on the Wireframe dataset (Huang et al. 2018) and the YorkUrban dataset (Denis, Elder, and Estrada 2008). We compare our method with the state-of-the-art methods including LSD (Von Gioi et al. 2008), DWP (Huang et al. 2018), AFM (Xue et al. 2019), LGNN (Meng et al. 2020), TP-LSD (Huang et al. 2020b), LETR (Xu et al. 2021), L-CNN (Zhou, Qi, and Ma 2019) , HAWP (Xue et al. 2020) and M-LSD (Gu et al. 2022). All the methods are learning-based methods except the classical LSD. The comparison results are listed in Table 2. RANK-LETR outperforms existing Transformer-based and CNN-based methods in prediction accuracy, particularly on the Wireframe dataset. Additionally, our method can converge with fewer training epochs compared to previous Transformer-based models. We also show the accuracy curves for line segment detection over the time of training in Fig.4. Our method reaches a high level of accuracy after just 60 epochs, which also proves that our method has a fast and stable convergence.

Visual examples of line segment detection results of two Transformer-based methods including LETR and ours on outdoor and indoor scenes in the Wireframe dataset are shown in Fig. 3. It is shown that our method can produce more accurate and complete detection result, especially in the detection of short line segments and in areas with high line density.

To verify the effectiveness of components and find the influence of the hyperparameters in our proposed method, we conduct an ablation and parameter study of our network architecture. The results are summarized in Table 3. The ablation study is conducted on the Wireframe dataset and the sAP and sF results are reported.

We employed the combination of ResNet50 and Transformer Encoder and Decoder from the classical LETR method as our baseline approach. We constructed a foundational version of our modeling using a ResNet50 network and a 6-layer Deformable Transformer Encoder, where feature maps at resolutions of $128\times 128$, $64\times 64$, and $32\times 32$ are input from the feature extraction backbone into the Transformer network for processing. Based on this architecture, we can use matched predicting strategy in supervising and find it can significantly improve the detection performance. On this basis, we sequentially added a re-ranking module and a rotation augmentation module, both of which further enhanced performance. Additionally, we compared the results of line segment prediction across different resolution feature maps, concluding that line segment prediction on high-resolution feature maps can provides a higher accuracy.

Moreover, we vary the number of referring points and the hidden dim of the Transformer architecture to show the influence of these hyperparameters. We find the hidden dim may be appropriately set to more that $128$ since a smaller one will lead to a lower accuracy. Changing the number of referring points has little impact on detection performance. In the ablation studies above, except for the baseline method that required 240 epochs to ensure convergence, all other methods were trained for 120 epochs, demonstrating the improved convergence capabilities of our proposed strategy. The saliency maps generated from some feature maps are also presented in Fig.5 to exhibit the learning effect. We also demonstrate the effectiveness of the ranking loss in the Table 4. We found that using the ranking loss can improve the performance of the model when trained on the Wireframe dataset and tested on the YUD dataset. This indicates that the ranking loss can enhance the generalization performance of the model.

While our method generally achieves the highest detection accuracy, it has a limitation when the centroids of two line segments are very close. In such cases, our method can only predict one of the segments. Although we find that nearby grid points may predict the corresponding line segment instead of the grid where the centroid lies, this remains a theoretical limitation of our approach. Allowing each feature point to predict more than one, e.g., two line segments could potentially mitigate this issue, which we leave as future work.