[Paper Review] TensorMask: A Foundation for Dense Object Segmentation

Xinlei Chen, Ross Girshick+2 more|arXiv (Cornell University)|Mar 28, 2019

Advanced Neural Network ApplicationsComputer Science37 references36 citations

TL;DR

TensorMask introduces a dense sliding-window framework for instance segmentation using structured 4D tensors, achieving competitive results to Mask R-CNN and providing a new foundation for dense mask prediction.

ABSTRACT

Sliding-window object detectors that generate bounding-box object predictions over a dense, regular grid have advanced rapidly and proven popular. In contrast, modern instance segmentation approaches are dominated by methods that first detect object bounding boxes, and then crop and segment these regions, as popularized by Mask R-CNN. In this work, we investigate the paradigm of dense sliding-window instance segmentation, which is surprisingly under-explored. Our core observation is that this task is fundamentally different than other dense prediction tasks such as semantic segmentation or bounding-box object detection, as the output at every spatial location is itself a geometric structure with its own spatial dimensions. To formalize this, we treat dense instance segmentation as a prediction task over 4D tensors and present a general framework called TensorMask that explicitly captures this geometry and enables novel operators on 4D tensors. We demonstrate that the tensor view leads to large gains over baselines that ignore this structure, and leads to results comparable to Mask R-CNN. These promising results suggest that TensorMask can serve as a foundation for novel advances in dense mask prediction and a more complete understanding of the task. Code will be made available.

Motivation & Objective

• Bridge the gap in dense sliding-window instance segmentation by treating masks as 4D tensors with geometric structure.
• Develop and evaluate the TensorMask representations and operators to enable dense mask prediction.
• Demonstrate that dense sliding-window approaches can achieve competitive performance to region-based methods on COCO.

Proposed method

• Represent masks with structured 4D tensors shaped (V, U, H, W) to capture both object position and mask geometry.
• Propose natural and aligned tensor representations and coordinate transformations between them.
• Introduce a tensor bipyramid to handle multi-scale masks, combining high-resolution masks for large objects with coarse localization for small objects.
• Attach a mask prediction head and a classification head to a feature pyramid network, with shared weights across levels.
• Use a class-agnostic mask head and a focal loss-based classification head, with a mask-specific per-pixel binary cross-entropy loss and an IoU-like box loss for training.
• Evaluate various tensor representations (natural vs aligned) and upscaling strategies, culminating in a tensor bipyramid that outperforms baseline heads.

Experimental results

Research questions

• RQ1 Can dense sliding-window instance segmentation be effectively realized using a 4D tensor representation of masks?
• RQ2 Do representations that respect the geometric structure of masks (natural vs aligned) yield improvements over unstructured approaches?
• RQ3 Does a tensor bipyramid enable high-quality masks across scales without prohibitive computation?
• RQ4 How does TensorMask perform compared to Mask R-CNN on COCO in terms of mask AP and box AP?

Key findings

• TensorMask with tensor bipyramid significantly improves mask AP over baseline heads using a feature pyramid, achieving 34.0 AP with 15x15 masks and up to 5.1 AP improvement over the best baseline (TensorMask vs feature pyramid).
• Upscaling with aligned representation (especially for large λ) yields large gains, e.g., aligned heads outperform natural heads by up to 9.2 AP at λ=5.
• The tensor bipyramid yields a substantial AP boost (≈5.1 AP) over the best baseline head on a feature pyramid and delivers higher-quality masks for large objects (AP L rises by 7.7).
• The tensor bipyramid with two window sizes (15x15 and 11x11) provides additional gains (~1.2 AP).
• Compared to Mask R-CNN, TensorMask with ResNet-50/101-FPN and training schedule achieves competitive results on COCO test-dev, with specific masks and box AP comparable to strong baselines.