Source-linked AI summary
TensorMask: A Foundation for Dense Object Segmentation
Xinlei Chen, Ross Girshick, Kaiming He, Piotr Dollár
TL;DR
Dense sliding-window instance segmentation remains under-explored compared with sliding-window box detection and detect-then-segment methods. TensorMask formalizes the task with structured 4D tensors and reports results near Mask R-CNN, suggesting a foundation for further dense mask-prediction research.
Problem
Dense sliding-window instance segmentation lacks direct approaches analogous to successful sliding-window box detectors, despite its scientific interest.
Method
TensorMask represents masks as structured 4D tensors with geometric sub-tensors and develops network architectures and operators that use this structure.
Results
TensorMask establishes a dense sliding-window instance segmentation system with results near Mask R-CNN and similar quantitative and qualitative results in example comparisons.
Takeaways & Limitations
The tensor view provides a foundation for exploring novel architectures and future research in dense sliding-window instance segmentation.
Takeaways & Limitations
The implementation assumes one sample per unit at integer coordinates and constrains α to positive integers for coordinate alignment.
Abstract
from arXiv · showhide
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.
1. Introduction
TensorMask addresses the lack of dense sliding-window instance segmentation by representing masks as structured 4D geometric tensors. Its tensor-based framework and tensor bipyramid achieve results near Mask R-CNN while handling objects across scales and overlap.
- Motivation: Dense sliding-window instance segmentation has lagged behind direct sliding-window box detection, motivating TensorMask as a foundation for this under-explored task.Existing instance segmentation methods commonly refine detected bounding boxes, whereas direct dense mask prediction lacks analogous approaches.
- Representation: The framework develops tensor-based mask representations and network architectures that explicitly capture mask geometry rather than packing masks into an unstructured channel axis.This geometric representation enables novel operations on the relative-mask sub-tensor.
- Results: Examples with a ResNet-101-FPN backbone show TensorMask and Mask R-CNN producing quantitatively and qualitatively similar results.The comparison uses the same images as Mask R-CNN’s Fig. 6.
- Results: TensorMask establishes a dense sliding-window instance segmentation system with results near Mask R-CNN and qualitatively reasonable delineation of small, large, and overlapping objects.The network and training procedure closely follow RetinaNet, with a dense mask predictor extending its dense bounding-box predictor.
- Representation: TensorMask represents masks with structured 4D tensors whose spatial and relative-mask axes both have geometric meaning.The tensor shape is (V, U, H, W), with (H, W) encoding object position and (V, U) encoding relative mask position.
- Architecture: TensorMask uses a tensor bipyramid whose scale-indexed 4D tensors grow in opposite directions across image and mask sub-tensors.This design gives large objects higher-resolution masks with coarse localization and small objects lower-resolution masks with fine localization.
2. Related Work
Related work spans proposal classification, detect-then-segment systems, pixel labeling with clustering, and dense mask proposals. TensorMask targets the missing dense sliding-window instance segmentation setting while adding structured representations and parallel multiclass prediction.
- Classify mask proposals: Early instance segmentation methods generated object mask proposals and then classified those proposals.This classify-mask-proposals methodology was also used for related detection and semantic segmentation tasks.
- Detect then segment: The dominant modern paradigm detects objects with bounding boxes and then segments each object using the box as a guide, led by Mask R-CNN.Mask R-CNN-based approaches have dominated recent instance segmentation challenge leaderboards.
- Dense sliding-window methods: Dense sliding-window instance segmentation was missing despite the success of sliding-window and region-based methods for box detection.TensorMask is presented as the first dense sliding-window approach for instance segmentation.
- Label pixels then cluster: Pixel-labeling approaches assign each pixel a category and auxiliary information, then cluster pixels into object instances.These methods natively predict higher-resolution masks for larger objects and benefit from semantic-segmentation advances.
- Dense sliding-window methods: DeepMask and InstanceFCN generate class-agnostic mask proposals densely, whereas TensorMask supports more expressive architectures and parallel multiclass classification.The related methods can be expressed as class-agnostic TensorMask models, but TensorMask targets multiclass instance segmentation.
3. Tensor Representations for Masks
TensorMask represents dense sliding-window masks with structured 4D tensors whose spatial sub-tensors encode both object position and relative mask position. Natural and aligned representations, coordinate transformations, upsampling, and tensor bipyramids make this geometry usable in network architectures.
- TensorMask represents dense sliding-window masks as structured 4D tensors with shape (V, U, H, W), rather than packing masks into an unstructured channel axis.The (V, U) and (H, W) sub-tensors both have geometric meaning, enabling operations on relative mask position and object position.
- 3.1. Unit of Length: The unit σ of each spatial axis specifies how one tensor pixel maps to input-image distance, and σVU may differ from σHW.Without units, the same tensor shape can represent windows of different image sizes.
- 3.2. Natural Representation: In the natural representation, the (V, U) sub-tensor at (y, x) represents mask values at offsets (y + αv, x + αu) within a window centered at (y, x).Here α=σVU/σHW, and the tensor assumes one sample per unit at integer coordinates.
- 3.3. Aligned Representation: The aligned representation stores, at each pixel, mask values from all windows overlapping that pixel, preserving pixel-to-pixel alignment during convolution.It is related to InstanceFCN score maps, but nearest-neighbor interpolation makes those maps unaligned and severely degrades results experimentally.
- 3.4–3.5. Transformations: The align2nat and nat2align transformations allow natural and aligned tensors to be used interchangeably, subject to integer-coordinate constraints handled through unit scaling.The up align2nat operator upsamples the (V, U) domain before conversion, generating high-resolution masks without inflating preceding channel counts.
- 3.6. Tensor Bipyramid: A tensor bipyramid adapts mask-pixel counts to scale: larger masks use larger (V, U) sub-tensors while the number of predicted masks decreases, with each level retaining the same total size.Its levels are constructed from scale-indexed 4D tensors using the framework’s alignment operations.
4. TensorMask Architecture
TensorMask models generate masks and categories in sliding windows over FPN features, using 4D tensor heads and a tensor bipyramid to preserve mask resolution across scales. Training uses mask-driven assignment and per-task losses, while baseline and tensor-bipyramid variants provide progressively richer architectural choices.
- Mask Prediction Heads: TensorMask models attach mask and classification heads to convolutional backbone features, with optional box prediction analogous to sliding-window detectors.The mask head generates windowed masks, while classification predicts object categories; box prediction is optional.
- Mask Prediction Heads: Baseline heads project each FPN feature map through a 1×1 convolution and reshape it into a specified 4D tensor.The four baseline heads use natural or aligned representations and natural or aligned upscaling variants.
- Mask Prediction Heads: Baseline models use one selected head across FPN levels, producing tensors whose sliding-window and feature-map units satisfy σVU=σHW.This couples the output window geometry to each feature-map resolution.
- Tensor Bipyramid: The tensor bipyramid head converts all levels to fine-resolution feature maps, applies conv+reshape, and then uses swap align2nat; it is the most effective explored variant.The tensor bipyramid preserves high mask resolution despite changes in spatial feature resolution.
- Training: Mask-driven label assignment marks a window positive only when containment, centrality, and uniqueness conditions all hold; otherwise it is negative.The assignment differs from IoU-based box-detector rules and works with only one or two window sizes and a single 1:1 aspect ratio.
- Training: Training uses per-pixel binary cross-entropy for masks, focal loss for classification, and parameter-free ℓ1 loss for box regression.Foreground mask pixels receive weight 1.5, and mask loss is averaged over each window’s pixels.
5. Experiments
Experiments evaluate TensorMask on COCO and show that structured representations, aligned upscaling, and tensor bipyramids substantially improve mask prediction. The best models approach Mask R-CNN accuracy but incur higher inference cost.
- Evaluation: TensorMask models are evaluated on COCO using mask average precision (AP), with training on approximately 118k train2017 images and validation on 5k val2017 images.Final results are reported on test-dev.
- TensorMask Representations: Simple natural and aligned heads perform similarly without upscaling, differing by only 0.4 AP.The comparison uses V × U=15×15 masks.
- TensorMask Representations: Aligned upscaling improves AP by 9.2 points over the natural head at λ=5, a 48% relative gain.The aligned head preserves sharp masks even with large upscaling factors.
- TensorMask Representations: Bilinear interpolation improves AP over nearest-neighbor interpolation by 3.1 points, especially at large upscaling factors.The difference is visually pronounced for overlapping objects.
- Tensor Bipyramid: The tensor bipyramid improves AP by 5.1 points over the best feature pyramid model, using 15×15 masks at k=0 and 480×480 masks at k=5.Higher-resolution masks at large-object levels provide the reported benefit.
- Comparison with Mask R-CNN: The best TensorMask reaches 35.4 mask AP on COCO test-dev versus 36.8 for Mask R-CNN, while ResNet-101 TensorMask reaches 37.1 with a 1.2 AP gap.These results are reported as close performance between dense sliding-window and detect-then-segment systems.
- Comparison with Mask R-CNN: The best R-101-FPN TensorMask runs at 0.38s per image versus 0.09s for Mask R-CNN because dense prediction evaluates more than 100k windows.Possible accelerations are outside the work’s scope.
A.1. Generalized Coordinate Transformation
The appendix generalizes the coordinate transformation between natural and aligned TensorMask representations by allowing arbitrary units for spatial and relative-mask axes.
- Generalized Transformation: The generalized transformation retains equal image-spatial units while allowing arbitrary units for image and relative-mask coordinates.The relations are stated in the image pixel domain, where the units are meaningful.
- Generalized Transformation: The generalized coordinate transformation demonstrates how units determine the mapping between aligned and natural representations.The authors suggest this formulation may enable more general uses.
- Generalized Transformation: Natural and aligned representations can be used interchangeably within one network through the coordinate transformation.This provides flexibility for designing network architectures.
A.2. Aligned Representation and InstanceFCN
The appendix formalizes the relationship between TensorMask’s aligned representation and InstanceFCN’s channel-based masks, showing that InstanceFCN corresponds to nearest-neighbor aligned prediction.
- InstanceFCN Correspondence: InstanceFCN’s output behaves as an upscaling aligned head using nearest-neighbor interpolation.This establishes a direct correspondence between the two representations.
- InstanceFCN Correspondence: InstanceFCN predicts a channel-indexed tensor whose channels correspond to spatial bins within each output mask.Rounding operations assign mask pixels to bins.
- InstanceFCN Correspondence: Within TensorMask notation, InstanceFCN masks are obtained by interpolating G on the relative-mask coordinates and applying align2nat with α=1.The construction converts the aligned result into natural masks.
- InstanceFCN Correspondence: TensorMask’s aligned representation differs because it gives geometric meaning to relative-mask axes, whereas InstanceFCN’s nearest-neighbor interpolation is unaligned.The paper reports that this unaligned behavior severely degrades results.
A.3. Object Detection Results
TensorMask also provides competitive bounding-box detection results, with box AP comparable to Mask R-CNN and better than RetinaNet under the reported setup.
- Object Detection Results: TensorMask has comparable bounding-box AP to Mask R-CNN and outperforms RetinaNet.The results are reported on COCO test-dev.
- Object Detection Results: The box-only TensorMask model resembles RetinaNet but uses mask-driven assignment and two window sizes instead of nine.All compared models use a ResNet-50-FPN backbone.
A.4. Mask-Only TensorMask
TensorMask can predict masks without a box head, while retaining nearly the same mask AP after accounting for the NMS setting. This suggests the box head is auxiliary in TensorMask, unlike in Mask R-CNN.
- Mask-Only TensorMask: TensorMask masks are independent of boxes, enabling mask-specific designs and experiments without a box head.Dense mask prediction still requires NMS; without regressed boxes, mask bounding boxes substitute for NMS.
- Mask-Only TensorMask: Switching from original boxes to mask bounding boxes for NMS causes a slight degradation in performance.
- Mask-Only TensorMask: After accounting for the NMS change, TensorMask without a box head has nearly equal mask AP to the mask+box variant.
- Mask-Only TensorMask: The box head is auxiliary in TensorMask, contrasting with its role in Mask R-CNN.
A.5. Qualitative Comparisons and Calibration
The qualitative comparisons use calibrated detection scores to display selected detections, and compare Mask R-CNN with TensorMask visually and by AP. TensorMask produces sharper-looking masks in the shown examples despite a lower AP.
- Qualitative Comparisons and Calibration: Visualizations display detections with a calibrated score ≥0.6.Calibration maps detector-category scores to precision values using precision-recall curves computed on val2017.
- Qualitative Comparisons and Calibration: With ResNet-101-FPN, Mask R-CNN obtains 38.3 AP and TensorMask obtains 37.1 AP on test-dev.
- Qualitative Comparisons and Calibration: TensorMask gives sharper masks visually than Mask R-CNN although its AP is 1 point lower.
- Qualitative Comparisons and Calibration: Figures 10 and 11 show continued qualitative results with Mask R-CNN in the top row and TensorMask in the bottom row of each set.