PaperlayerView exact revision on arXiv

Source-linked AI summary

ISTR: End-to-End Instance Segmentation with Transformers

Jie Hu, Liujuan Cao, Yao Lu, ShengChuan Zhang, Yan Wang, Ke Li, Feiyue Huang, Ling Shao, Rongrong Ji

arXiv:2105.00637v2cs.CV

TL;DR

Instance segmentation has lacked an end-to-end approach comparable to those used in object detection, in part because masks are higher-dimensional and matched training samples are sparse. ISTR predicts low-dimensional mask embeddings and recurrently refines detection and segmentation together. On COCO, it reports 46.8/38.6 box/mask AP with ResNet50-FPN and 48.1/39.9 box/mask AP with ResNet101-FPN.

  • Problem

    End-to-end set prediction has been difficult to apply to instance segmentation, where high-dimensional masks and sparse matched samples make direct mask-head learning challenging.

    Source: §1 · p. 1

  • Method

    ISTR predicts low-dimensional mask embeddings for set-loss matching and uses recurrent refinement to update predictions for detection and segmentation together.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    ISTR reports 46.8/38.6 box/mask AP with ResNet50-FPN and 48.1/39.9 box/mask AP with ResNet101-FPN on COCO test-dev.

    Source: §1 · p. 2

  • Takeaways & Limitations

    ISTR's COCO performance demonstrates its potential as a baseline for instance-level recognition.

    Source: §5 · p. 9

    Although ISTR achieves state-of-the-art mask AP, its masks could be improved by learning finer masks.

    Source: §4.3 · p. 8

Abstract

from arXiv · show

End-to-end paradigms significantly improve the accuracy of various deep-learning-based computer vision models. To this end, tasks like object detection have been upgraded by replacing non-end-to-end components, such as removing non-maximum suppression by training with a set loss based on bipartite matching. However, such an upgrade is not applicable to instance segmentation, due to its significantly higher output dimensions compared to object detection. In this paper, we propose an instance segmentation Transformer, termed ISTR, which is the first end-to-end framework of its kind. ISTR predicts low-dimensional mask embeddings, and matches them with ground truth mask embeddings for the set loss. Besides, ISTR concurrently conducts detection and segmentation with a recurrent refinement strategy, which provides a new way to achieve instance segmentation compared to the existing top-down and bottom-up frameworks. Benefiting from the proposed end-to-end mechanism, ISTR demonstrates state-of-the-art performance even with approximation-based suboptimal embeddings. Specifically, ISTR obtains a 46.8/38.6 box/mask AP using ResNet50-FPN, and a 48.1/39.9 box/mask AP using ResNet101-FPN, on the MS COCO dataset. Quantitative and qualitative results reveal the promising potential of ISTR as a solid baseline for instance-level recognition. Code has been made available at: https://github.com/hujiecpp/ISTR.

1. Introduction

ISTR addresses the difficulty of making instance segmentation end-to-end by predicting low-dimensional mask embeddings and refining detection and segmentation together. Its COCO results show strong performance with ResNet50-FPN and ResNet101-FPN.

1. Introduction

  • End-to-end instance segmentation remains challenging because prior approaches use NMS or decompose the task, while direct set-loss adaptation performs poorly with sparse mask-head training samples.Masks can be 28 × 28 or larger, whereas COCO provides an average of 7.7 matched ground-truth boxes per image, compared with 512 proposals used by Mask R-CNN to train its mask head.Source: §1 · p. 1
  • ISTR predicts low-dimensional mask embeddings for set-loss matching, replacing direct regression of high-dimensional masks.The paper reports that the first few mask components contain most mask information, motivating embeddings that can be learned from fewer matched samples.Source: §1 · p. 2, Figure 1 · p. 2
  • 46.8/38.6 box/mask AP with ResNet50-FPN and 48.1/39.9 box/mask AP with ResNet101-FPN demonstrate ISTR's COCO performance.The results use the test-dev split and were measured at 13.8 fps and 11.0 fps, respectively, with a single 1080Ti GPU.Source: §1 · p. 2
  • The paper presents ISTR as the first Transformer-based end-to-end instance-segmentation framework and reports accuracy and runtime on par with state-of-the-art methods on COCO.The approach uses approximation-based suboptimal embeddings obtained through PCA's closed-form solution.Source: §1 · p. 2
  • ISTR recurrently refines predictions to detect and segment instances concurrently, offering an alternative to top-down and bottom-up frameworks.The paper presents this strategy as a way to process the two tasks together.Source: §1 · p. 2

2. Related Work

Prior instance-segmentation methods use top-down, bottom-up, or box-independent approaches, while earlier end-to-end trials were limited in evaluation. ISTR predicts boxes and mask embeddings together, and jointly processes detection and segmentation.

2. Related Work

  • Top-down methods detect then segment, bottom-up methods classify pixels then cluster them, and SOLO avoids dependence on box detection.ISTR instead predicts sets of bounding boxes and mask embeddings, avoiding decomposition into dependent tasks; the paper contrasts this with MEInst's redundant predictions and suboptimal performance.Source: §2 · p. 2
  • Earlier end-to-end instance-segmentation trials used recurrent neural networks but were evaluated only on small datasets and not against current baselines.ISTR uses mask-embedding similarity for bipartite matching and incorporates Transformers.Source: §2 · p. 2
  • ISTR targets single-image instance segmentation, unlike the cited Transformer approach for end-to-end video instance segmentation that uses continuous frames.Figure 2 describes its pipeline: image features and RoI features feed a Transformer encoder, followed by recurrent prediction refinement.Source: §2 · p. 3, Figure 2 · p. 3
  • ISTR processes detection and segmentation concurrently, consistent with prior work reporting benefits from jointly learning the tasks.The related-work discussion cites Mask R-CNN and other multitask mechanisms as examples of joint-learning studies.Source: §2 · p. 3

3. Proposed Method

ISTR represents masks with low-dimensional embeddings, uses bipartite matching and set loss to train predictions, and refines detection and segmentation together through recurrent updates.

3.2. Matching Cost and Prediction Loss

  • ISTR predicts low-dimensional mask embeddings and matches them with ground-truth embeddings using bipartite matching and a set prediction loss.The set loss regresses matched predictions against ground-truth targets.Source: §3 · p. 3, §3.2 · p. 4

3.1. Mask Embeddings

  • The mask-embedding objective becomes the PCA objective under linear encoder and decoder transformations, yielding a closed-form solution for the transformation.The objective can also be optimized with other models, such as an autoencoder.Source: §3.1 · pp. 3–4

3.3. Instance Segmentation Transformer

  • ISTR combines a CNN-FPN backbone, dynamic-attention Transformer encoder, concurrent prediction heads, and recurrent updates to refine detections and segmentations.The encoder learns object relations, while each update refines the query boxes and prediction set.Source: §3.3 · pp. 4–5
  • Replacing masks with mask embeddings improves mask AP; performance improves at embedding dimension l=60, saturates at l=80, and is worse when masks are directly expanded as embeddings with l=784.The ablations train on COCO train2017 and report mask and box AP on val2017.Source: §3.3 · p. 5, Table 1 · p. 5
  • Using both pixel-level dice loss and embedding-level L2 loss yields gains in mask AP.Source: §3.3 · p. 5
  • Cosine similarity as the mask matching cost brings expected gains, while mask dice-cost matching performs slightly better than omitting a mask cost.L1 matching between predicted and encoded embeddings also gives a slight improvement.Source: §3.3 · p. 5
  • Dynamic attention significantly outperforms multi-head attention for fusing RoI and image features; global average pooling beats max-pooling, and position embeddings improve box and mask AP.Source: §3.3 · p. 5

4. Experiments

Experiments on COCO evaluate ISTR’s mask representations, matching and loss choices, attention and pooling, and recurrent refinement. On the test-dev split, ISTR reports gains over instance-segmentation and end-to-end detection baselines, while qualitative results identify finer masks as an area for future work.

4. Experiments

  • On MS COCO, models are trained on train2017 (118K images), ablations are tested on val2017 (5K images), and final results are reported on test-dev.Final models use a 36-epoch schedule, with learning-rate drops at epochs 27 and 33.Source: §4 · pp. 5–6

4.3. Main Results

  • With ResNet101-FPN, ISTR exceeds MEInst by 4.6 points in APm and DETR and sparse R-CNN by 4.6 and 2.5 points in APb.Table 2 reports quantitative results on COCO test-dev; all models use multi-scale training.Source: Table 2 · p. 7, §4.3 · p. 8

4.1. Ablation Experiments

  • Mask embeddings outperform raw masks, improve as dimension l reaches 60, saturate at l=80, and perform worse when expanded to l=784.The authors attribute the advantage over high-dimensional masks to embeddings being more effectively regressed with few matched samples.Source: §4.1 · p. 7
  • Cosine similarity between embeddings improves mask matching, while combining pixel-level dice and embedding-level L2 losses performs better than either loss alone.Embedding-level L2 improves performance, but the learned embeddings are slightly suboptimal for reconstructing the original masks.Source: §4.1 · p. 7
  • Dynamic attention outperforms multi-head attention, and global average pooling with summed position embeddings performs better than max pooling.The authors suggest a single projection is more effective for learning relations between objects; they describe position embeddings as essential for high results.Source: §4.1 · p. 7

4.2. Stage Analysis

  • Across six stages, boxes and masks refine step by step, while performance saturates in the last two stages.The qualitative and quantitative results show progressive refinement.Source: Figure 3 · p. 6, §4.2 · p. 8

4.3. Main Results

  • Qualitative comparisons show Mask R-CNN retains duplicate predictions without NMS, and ISTR’s masks could be made finer.The authors leave learning finer masks for future work.Source: §4.3 · p. 8, Figure 4 · p. 8

5. Conclusion

ISTR uses low-dimensional mask embeddings and recurrent refinement to conduct detection and segmentation concurrently. Its COCO performance demonstrates potential for instance-level recognition.

5. Conclusion

  • ISTR predicts low-dimensional mask embeddings for matching and regression, and conducts detection and segmentation concurrently through recurrent refinement.The conclusion describes the framework as a new way to explore end-to-end instance segmentation.Source: §5 · p. 9
Loading 2105.00637v2…