Source-linked AI summary
ShapeMask: Learning to Segment Novel Objects by Refining Shape Priors
Weicheng Kuo, Anelia Angelova, Jitendra Malik, Tsung-Yi Lin
TL;DR
Precise mask annotations for every category limit instance segmentation of novel categories. ShapeMask refines box detections with shape priors and instance embeddings; it outperforms MaskX R-CNN in cross-category transfers and is competitive in fully supervised segmentation.
Problem
Precise pixelwise supervision for every category limits instance segmentation to a small portion of the visual world with dense annotations, leaving generalization to novel categories an open challenge.
Source: §1 · p. 1Method
ShapeMask refines box detections by selecting shape priors for an object's rough shape, then decodes and refines the mask using instance embeddings.
Source: §1 · p. 2Results
ShapeMask outperforms MaskX R-CNN by 6.4 AP on VOC-to-Non-VOC transfer and 3.8 AP on Non-VOC-to-VOC transfer.
Source: §5.1 · p. 6Takeaways & Limitations
ShapeMask is robust to inaccurate detections, competitive in fully supervised segmentation, and efficient for training and inference.
Source: §6 · p. 8
Abstract
from arXiv · showhide
Instance segmentation aims to detect and segment individual objects in a scene. Most existing methods rely on precise mask annotations of every category. However, it is difficult and costly to segment objects in novel categories because a large number of mask annotations is required. We introduce ShapeMask, which learns the intermediate concept of object shape to address the problem of generalization in instance segmentation to novel categories. ShapeMask starts with a bounding box detection and gradually refines it by first estimating the shape of the detected object through a collection of shape priors. Next, ShapeMask refines the coarse shape into an instance level mask by learning instance embeddings. The shape priors provide a strong cue for object-like prediction, and the instance embeddings model the instance specific appearance information. ShapeMask significantly outperforms the state-of-the-art by 6.4 and 3.8 AP when learning across categories, and obtains competitive performance in the fully supervised setting. It is also robust to inaccurate detections, decreased model capacity, and small training data. Moreover, it runs efficiently with 150ms inference time and trains within 11 hours on TPUs. With a larger backbone model, ShapeMask increases the gap with state-of-the-art to 9.4 and 6.2 AP across categories. Code will be released.
1. Introduction
ShapeMask addresses novel-category instance segmentation by refining box detections through shape priors and instance embeddings. It outperforms the state-of-the-art cross-category approach, remains competitive when fully supervised, and runs efficiently.
1. Introduction
- ShapeMask is designed to start from approximate localization and progressively refine the shape into a final mask.The bounding box need not be accurate.Source: Figure 1 · p. 1
- Precise pixelwise supervision for every category limits instance segmentation to a small portion of the visual world with dense annotations.The introduction frames generalization to novel categories as an open challenge.Source: §1 · p. 1
- Shape priors offer richer shape information than box priors, while grouping-based instance embeddings can use appearance similarity to associate pixels with an object.The authors hypothesize that shared shapes across categories may aid generalization when novel-object pixels look different.Source: §1 · p. 2, Figure 2 · p. 2
- ShapeMask first selects shape priors to capture an object's location, scale, and rough shape, then decodes and refines a coarse mask with instance embeddings.It uses class-agnostic training for novel categories without transfer learning; its soft shape priors can also produce masks beyond the detection box, unlike feature-cropping methods.Source: §1 · p. 2
- Cropping features instead of using ROIAlign, jittering ground truths, and using RetinaNet enable training 4× faster without compromising performance.These choices are designed for efficient operation across TPUs and GPUs.Source: §1 · p. 2
- ShapeMask significantly outperforms the state-of-the-art transfer-learning approach in cross-category experiments and is competitive in fully supervised segmentation while running at 150–200 ms per image.The paper also reports training multiple times faster than state-of-the-art techniques.Source: §1 · p. 2
2. Related Work
Prior work divides instance segmentation into detection-based methods that predict masks within boxes and grouping-based methods that assemble instances from pixelwise cues. Related approaches use shape priors, while novel-category methods and partially supervised setups address missing mask labels in different ways.
2. Related Work
- Detection-based methods predict masks within detected boxes, whereas grouping-based methods group instances from pixelwise cues such as affinity, vectors, watershed energy, or semantic classes.Grouping methods typically perform post-processing to assemble instances.Source: §2 · p. 2
- Shape priors have been used as probabilistic unaries, augmented proposals, or top-down cues for grouping, but grouping methods had not surpassed detection-based methods on COCO.Classical grouping approaches showed minimal generalization gaps to unseen categories.Source: §2 · p. 3
- Novel-category methods include pseudo-mask supervision from GrabCut, boundary detection followed by grouping, and image-level supervision, but do not exploit existing instance-mask labels.These works address categories without mask annotations.Source: §2 · p. 3
- Partially supervised instance segmentation trains with box labels for some categories and requires masks at test time; MaskX R-CNN predicts mask-branch weights from the box-detection branch.MaskX R-CNN improves on class-agnostic training but retains a gap from fully supervised systems.Source: §2 · p. 3
3. Method
ShapeMask successively refines box detections into instance masks through shape recognition, coarse mask prediction, and instance-specific refinement. Its class-agnostic setup uses shape priors learned across categories to support segmentation of novel objects.
3.1. Shape Recognition
- Shape priors are learned by clustering training masks into centroids, with K priors per class-agnostic model to capture diverse shapes across categories.Masks are resized to 32 x 32 before clustering.Source: §3.1 · p. 3
- A pooled box-feature embedding selects and combines similar shape priors into a predicted shape, which is fitted to the detection box as a smooth detection prior.Pixel-wise MSE against the ground-truth mask trains the prior prediction; this shape-first step regularizes outputs against implausible shapes and is associated with improved generalization to novel categories.Source: §3.1 · pp. 3–4
3.2. Coarse Mask Prediction
- The detection prior is embedded and combined with image features to form prior-conditioned features, which a convolutional decoder maps to a coarse instance mask.The coarse mask is trained with pixel-wise cross-entropy.Source: §3.2 · p. 4
- Because the coarse mask lacks instance-specific image information, ShapeMask pools features within it, subtracts that embedding from prior-conditioned features, and decodes a higher-resolution fine mask.The fine mask is trained against the higher-resolution ground-truth mask.Source: §3.3 · p. 4
3.4. Generalization by Class Agnostic Learning
- For novel categories, the mask branch predicts a foreground mask without class identity, using shape priors clustered across all classes and treating novel objects as foreground at inference.The class-agnostic model uses a larger K to capture diverse shape modes, while ShapeMask remains competitive in the class-specific setting.Source: §3.4 · pp. 4–5
4. Implementation Details
ShapeMask’s implementation trains mask prediction from RetinaNet detections and jittered ground-truth boxes, then processes box-centered features through a feature pyramid. Its patches are sized to contain the instance without crop-and-resize.
4. Implementation Details
- ShapeMask trains the mask branch on a fixed number of ground-truth masks and boxes per image, rather than sampling masks from object proposals.The sampled boxes are jittered to better mimic imperfect inference detections.Source: §4 · p. 5
- Gaussian-jittered boxes become uniform box priors, and jittering helps the mask branch learn robustness to noisy detections at test time.The box centers and sizes are perturbed with Gaussian noise; the stated standard deviation is 0.1.Source: §4 · p. 5
- RoIs use P3–P5 pyramid levels and centered feature patches sized to contain the entire instance, without a crop-and-resize operation.A 1×1 convolution reduces patch features from 256 to 128 dimensions.Source: §4 · p. 5
5. Experiments
The experiments evaluate ShapeMask on novel-category transfer, fully supervised COCO segmentation, and component ablations, alongside training and inference efficiency.
5.1. Generalization to Novel Categories
- ShapeMask beats MaskX R-CNN by 6.4 AP on VOC→Non-VOC and 3.8 AP on Non-VOC→VOC; with a stronger feature pyramid, margins rise to 9.4 and 6.2 AP.The comparisons use the same ResNet backbone for the first two transfers.Source: Table 1 · p. 6, §5.1 · p. 6
- With 1/100 training data, ShapeMask beats full-data MaskX R-CNN by 2.0 AP; COCO-trained models also segment robotics objects without domain fine-tuning.The robotics evaluation uses ground-truth boxes and data without instance-mask annotations.Source: Figure 8 · p. 7, §5.1 · p. 7, Figure 9 · p. 7
5.2. Fully Supervised Instance Segmentation
- Training takes 11 hours on TPUs, 4x faster than Mask R-CNN; mask-branch capacity falls 130x with a 6x speedup (4.6ms) at marginal performance loss.For heavier models, training finishes within 25 hours on TPUs and inference runs at 5 fps per 1024 x 1024 image on GPU.Source: §5.2 · pp. 7–8
- On fully supervised COCO, ShapeMask beats Mask R-CNN by 1.7 AP with ResNet-101-FPN; a larger backbone yields gains of 2.9 AP over Mask R-CNN and 2.7 AP over MaskLab.ShapeMask remains 2.0 AP behind PANet without techniques reported in [33, 5].Source: §5.2 · p. 7, Table 2 · p. 8
5.3. Model Ablation
- Shape priors and instance embeddings independently improve the partially supervised baseline by about 12 and 5 AP, respectively, and combining them improves further.Similar ablation results are reported for the fully supervised setting.Source: §5.3 · p. 8, Table 4 · p. 8
6. Conclusion
ShapeMask uses shape priors and instance embeddings to improve generalization to novel categories, while remaining competitive in fully supervised segmentation. The conclusion highlights its ability to handle poorly localized detections and produce quality contours, including on thin structures and overlapping objects.
6. Conclusion
- ShapeMask significantly outperforms the state of the art in cross-category segmentation and is competitive in the fully supervised setting.The conclusion also characterizes it as efficient for training and inference.Source: §6 · p. 8
- Visualizations show quality contours for large objects, thin structures, and within-category overlaps.Source: Figure 10 · p. 9
- ShapeMask handles poorly localized detections at test time, whereas Mask R-CNN fails under the same stress test by design of tight feature cropping.Both models were trained on whole boxes; ShapeMask uses soft detection priors.Source: Figure 11 · p. 9
Appendix A: ShapeMask vs. Mask R-CNN
The appendix compares partially supervised ShapeMask and Mask R-CNN on transfers between VOC and non-VOC categories. Both segment unseen categories, while ShapeMask provides more accurate segments; both have challenges on pizza and people.
Appendix A: ShapeMask vs. Mask R-CNN
- ShapeMask provides more accurate segments than Mask R-CNN in visualized transfers between VOC and non-VOC categories.The images were randomly sampled from the validation set and show both VOC-to-non-VOC and non-VOC-to-VOC transfers.Source: Appendix A · p. 9
- Both methods segment unseen object categories in the displayed transfers.Source: Appendix A · p. 9
- Pizza and people are challenging categories where both methods exhibit potential failure modes.Source: Appendix A · p. 9
Appendix B: Object Detection
On COCO object detection, ShapeMask performs competitively against leading detectors, with results depending on the backbone. With ResNet-101-FPN it outperforms RetinaNet and Mask R-CNN; with ResNet-101-NAS-FPN it is comparable to SNIP and trails PANet.
Appendix B: Object Detection
- 42.0 AP with ResNet-101-FPN outperforms RetinaNet and Mask R-CNN and is among the best reported approaches using that backbone.Source: Appendix B · p. 9
- With ResNet-101-NAS-FPN, ShapeMask achieves 45.4 AP, comparable to SNIP and 2.0 AP behind PANet.Source: Appendix B · p. 9
- ShapeMask functions as a competitive object detector without applying many existing detection improvement methods.The cited methods are [4, 21, 44].Source: Appendix B · p. 9
Appendix C: Lightweight Mask Branch
ShapeMask’s mask branch retains competitive performance as its capacity falls substantially. At 16 channels, it uses far fewer parameters and FLOPs than Mask R-CNN, with mask-branch timing reported separately.
Appendix C: Lightweight Mask Branch
- ShapeMask performance decreases only slightly as mask-branch capacity drops dramatically.The branch’s convolution and deconvolution layers use the same channel count in this comparison.Source: Appendix C · p. 9, Table 6 · p. 10
- The 16-channel mask branch runs at 4.6ms; the reported timing measures the mask branch only.Source: Appendix C · p. 9, Table 6 · p. 10
- With 16 channels, ShapeMask is 0.4 AP higher than Mask R-CNN while using 130x fewer parameters and 23x fewer FLOPs.Source: Appendix C · p. 9, Table 6 · p. 10
Appendix D: Ablation Studies of Fully Supervised ShapeMask
In the fully supervised COCO ablation, shape priors and instance embeddings each improve the baseline, and using both boosts performance further.
Appendix D: Ablation Studies of Fully Supervised ShapeMask
- Using both shape priors and instance embeddings boosts performance further than either component alone.The ablation uses COCO val2017 with ResNet-101-FPN.Source: Appendix D · p. 10
- Using an object shape prior improves performance over the baseline.Source: Appendix D · p. 10
- Using an instance embedding improves performance over the baseline.Source: Appendix D · p. 10