Amodal Instance Segmentation with Diffusion Shape Prior Estimation

Authors: minh  tran
khoa  vo
vuong  ho
tri  nguyen
ngan  le
Venue: thecvf_com_CVPR_2024_Workshop_POETS
Published: 2024-04-16 11:51:49

Abstract: Amodal Instance Segmentation (AIS) presents an intriguing challenge, including the segmentation prediction of both visible and occluded parts of objects within images. Previous methods have often relied on shape prior information gleaned from training data to enhance amodal segmentation. However, these approaches are susceptible to overfitting and  disregard object category details. Recent advancements highlight the potential of conditioned diffusion models, pretrained on extensive datasets, to generate images from latent space. Drawing inspiration from this, we propose AISDiff with a Diffusion Shape Prior Estimation (DiffSP) module. AISDiff begins with the prediction of the visible segmentation mask and object category, alongside occlusion-aware processing through the prediction of occluding masks. Subsequently, these elements are inputted into our DiffSP module to infer the shape prior of the object. DiffSP utilizes conditioned diffusion models pretrained on extensive datasets to extract rich visual features for shape prior estimation. Additionally, we introduce the Shape Prior Amodal Predictor, which utilizes attention-based feature maps from the shape prior to refine amodal segmentation. Experiments across various AIS benchmarks demonstrate the effectiveness of our AISDiff.

Main paper: 
 introduction: 
 Amodal perception, as described in [10] , describe human's remarkable ability to perceive objects in their entirety despite occlusion. Building upon this concept, the pioneering studies by [12, 30] introduced amodal instance segmentation (AIS). This approach aims to predict the complete shape of objects, encompassing both their visible and occluded regions. Indeed, AIS exhibits vast potential across various domains, as evidenced by its applications in robot manipulation [1] and autonomous driving [19] . Across various AIS benchmarks [3, 19, 30] , a multitude of approaches addressing the AIS challenge have emerged in the literature. These approaches, as evidenced by numerous studies [3, 7, 12, 14, 19, 24, 25] , demonstrate the ongoing efforts to Next, these predictions are integrated into the Diffusion Shape Prior Estimation (DiffSP) module to establish the object's shape prior. This shape prior is then utilized by AISDiff to produce the amodal segmentation.tackle this challenge.Recent research [2, 4, 7, 26, 27] highlights the effectiveness of integrating shape prior information in AIS. Indeed, These shape prior AIS methods typically construct shape-prior knowledge from the training dataset, which is later utilized to train the AIS model. In [26] , for instance, the authors employ variational autoencoders to reconstruct amodal masks. The concept revolves around using ground truth amodal masks, utilizing autoencoders to reconstruct them, and storing the encoded codebooks as shape priors. Similarly, in [4] , the authors also construct a shape prior codebook but employ a vector-quantization variational au-toencoder. After establishing the shape prior, these method first predict the coarse amodal segmentation and refine the final amodal segmentation mask using the built shape prior. However, there are limitations to these approaches. Firstly, the shape prior tends to overfit to the training data, consequently leading to overfitting in amodal mask prediction overall. Secondly, since the shape prior is built solely from ground truth amodal masks, it may overlook the object category, which could provide significant supplementary information for deriving the shape prior.To tackle these issues, we desgin a AIS mask head with Diffusion Shape Prior Estimation (AISDiff). The design of AISDiff is depicted in Figure 1 . In essence, AISDiff begins by predicting the visible segmentation mask and the category of the object of interest. Simultaneously, it conducts occlusion-aware processing by predicting the occluding mask, which is the segmentation of occluding elements within the specified ROI. Subsequently, these three pieces of information are fed into the proposed Diffusion Shape Prior Estimation (DiffSP) module to derive the shape prior of the object. Finally, leveraging this shape prior, AISDiff generates the amodal segmentation.Specifically, DiffSP leverages the successes of conditioned diffusion models (such as Stable Diffusion [20] and GLIDE [17] ), which are pretrained on extensive language vision datasets like LAION [22] . This enables the model to capture rich visual features, making it suitable as prior knowledge for downstream tasks [18, 28] . Building upon this foundation, we feed a trained conditioned diffusion model with an ROI image containing only the visible pixels of the object of interest, expecting the model to generate the missing parts. Additionally, an occluding mask and a textual description of the object category is also feed to condition the mdoel. Subsequently, the denoising process iterates T steps to output the generated image containing the occluded parts. However, rather than relying on the final generated pixels, DiffSP exploits on the attention mechanism between the conditioning information and the image features. This attention map remains relatively stable across time steps, thereby reducing the denoising time needed to obtain the shape prior. Furthermore, we design the Shape Prior Amodal Predictor, which learns the attention-based amodal feature map from the acquired shape prior to predict the amodal mask segmentation.In summary, our contributions are as follows: • We present AISDiff, a novel AIS mask head featuring a Diffusion Shape Prior Estimation module. This model predicts the visible segmentation mask and category of the object while considering occlusion. It then uses these predictions to estimate the shape prior of the object before generating the final amodal segmentation mask. • We propose DiffSP module, harnessing the efficacy of conditioned diffusion models to derive the shape prior of the object of interest. • We introduce the Shape Prior Amodal Predictor, which learns attention-based amodal feature maps from the obtained shape prior to predict the amodal segmentation. 
related work: 
 Amodal instance segmentation involves predicting an object's shape, including both its visible and occluded parts. Li and Malik [12] pioneered a method aimed at addressing AIS. They proposed enlarging the modal bounding box in alignment with high heatmap values and synthesizing occlusions. Following this seminal work, various methodologies have surfaced in literature. Notably, ORCNN [3] introduces instance mask heads for both amodal and visible instances, along with an additional head for predicting occluded masks. ASN [19] builds upon ORCNN by integrating a multi-level coding module for bidirectional feature modeling of visible and amodal aspects. BCNet [8] enhances amodal mask prediction by incorporating a supplementary branch dedicated to predicting occlusion masks within the bounding box. AISFormer [24] introduces a transformer-based mask head, demonstrating the efficacy of transformer modeling in generating AIS masks. However, their approach, which consolidates all mask relationships into one transformer model, leads to compromised visible segmentation output, consequently affecting the quality of amodal segmentation output due to bidirectional feature relations as mentioned earlier.Recent studies [7, 26] underscore the benefits of integrating shape priors into AIS. These methods leverage prior knowledge of mask shapes to improve amodal mask predictions. VRSP-Net [26] predicts coarse amodal masks, retrieves shape priors using a simple autoencoder, and then refines the final amodal mask predictions. AmodalBlastomere [7] employs a similar strategy with a variational autoencoder for blastomere and cell segmentation. Despite their progress, these methods often overlook the importance of object categories when utilizing prior shapes. Moreover, their training procedures frequently lead to overfitting of the shape prior model to the training dataset. Additionally, these approaches simply incorporate the shape prior by concatenating it with visible features to refine amodal masks. 
overall ais setup: 
 Given an input image I, we follow most of previous AIS settings [3, 8, 24, 26] , utilizing a pre-trained backbone network, such as ResNet [6] , RegNet [21] to extract spatial visual representation. An object detector such as FCOS [23] , or Faster-RCNN [6] , can be subsequently adopted to obtain n regions of interest (RoI) predictions and their corresponding visual features {F i } n i=1 . Follow most of previous works [8, 24, 26], the object detector being chose is Faster R-CNN for fair comparison. Here, each RoI is presented by its visual feature F i ∈ R Ce×Hr×Wr , where C e denotes the feature channel size and H r × W r represents the spatial shape of the pooling feature. In this context, given a RoI, AISDiff takes F i as input and aims to predict the amodal mask M i a . Moreover, in this case, we also denote the visible mask M i v , and the occluding mask M i o .The overall design of AISDiff is depicted in Fig. 1 . Initially, we discuss the prediction process for the visible segmentation of the object of interest, along with its categories, incorporating occlusion-awareness through the prediction of occluding masks (Sec. 3.2.1). Following this, we introduce the DiffSP method in detail (Sec. 3.2.2). Lastly, we present the Shape Prior Amodal Predictor (Sec. 3.2.3).Given the ROI feature F i , AISDiff first aims to predict the visible segmentation mask and the category of the object of interest, while simultaneously conducts occlusion-aware ability by predicting the occluding mask, which is the segmentation of occluding elements within the specified ROI.BCNet [8] is utilized as the foundation for the Occlusionaware Visible Segmentation module. This module consists of two branches: one for occluding mask prediction and the other for visible mask prediction. Drawing from the methodology outlined in [8] , both branches follow a similar design structure, encompassing two main components: feature extraction and mask prediction. The feature extraction segment comprises a sequence of layers, including a 3 × 3 convolutional layer with a stride of 1, a Graph Convolutional Network (GCN) [11] block, and another 3 × 3 convolutional layer with a stride of 1. Subsequently, the mask prediction component is constructed with a 2×2 transposed convolutional layer employing a stride of 2, coupled with a 1 × 1 convolutional layer using a stride of 1. Furthermore, to enhance occlusion awareness and subsequently improve visible segmentation accuracy, features extracted from the occluding branch are incorporated into the ROI feature F i before being fed into the feature extraction section of the visible branch. Simultaneously, features extracted from the visible branch are utilized for object category prediction. This classification step employs a fully connected layer with an output dimension corresponding to the number of categories present in the datasets under consideration. In summary, the final output of this module comprises the visible mask M i a , the occluding mask M i o , and the object category c i .The process depicted in Fig. 2 illustrates the Shape Prior Estimation (DiffSP) module. DiffSP builds upon the successes of conditioned diffusion models, such as Stable Diffusion [20] and GLIDE [17] , which are pre-trained on comprehensive language-vision datasets like LAION [22] . This pretraining equips the model with the ability to capture intricate visual features, rendering it suitable as prior knowledge for subsequent tasks [18, 28] . Expanding on this foundation, DiffSP utilizes a trained conditioned diffusion model and inputs a ROI image containing only the visible pixels of the object under consideration, expecting the model to generate the obscured parts. Additionally, the model is conditioned with an occluding mask and a textual description of the object category. Subsequently, the denoising process iterates T steps to produce the generated image containing the occluded regions. However, instead of relying solely on the final generated pixels, DiffSP capitalizes on the attention mechanism between the conditioning information and the image features.Specifically, Stable Diffusion [20] is employed as the pre-trained conditioned diffusion model, leveraging its self and cross-attention layers. Specifically, the random Gausian noise is encoded into latent space and then experiences the denoising process over T time steps to generate the inpainting image. In fact, the ROI image containing only the visible pixels of the object of interest, the occluding mask, and the textual description of the object category serve as conditions and are represented as y, which is projected by τ into an intermediate representation τ (y). At each denoising step t, a UNet architecture with L layers of self and crossattention transforms z t into z t-1 . Specifically, at layer l and time step t, the cross-attention layer captures the relationship between z t and the encoded condition τ (y), reflecting the entire reconstructed shape of the object. This relationship is formalized as follows: at layer l and time step t, the self-attention map is denoted as A l,t S , and the cross-attention map is denoted as A l,t C . Moreover, as demonstrated in [16] , the attention map remains relatively stable across time steps. Following the methodology of [16] , we average these cross and self-attention maps over layers and time steps, setting T = 10. Additionally, as also suggested in [16] , although the cross-attention maps A C already outline the shape of the reconstructed object, they tend to be coarse-grained and noisy. To refine the precision of object localization, we follow [16] , utilizing the self-attention map A S to enhance A C . Consequently, the shape prior is obtained by: M sp = (A S ) τ • A C .The design of Shape Prior Amodal Predictor is depicted in Fig. 3 . Initially, the feature extraction module utilizes the ROI feature F i to generate the amodal feature. This module is constructed using a sequence of 3 × 3 convolutional layers with a stride of 1. Subsequently, the obtained amodal feature undergoes processing in the attention learning module in conjunction with the shape prior M sp obtained from DiffSP, aimed at learning the spatial attention map. Specifically, the attention computation involves passing the amodal feature through a sequence of 3 × 3 convolutional layers with a stride of 1, followed by a sigmoid activation function. This computed attention map is then multiplied with the shape prior M sp . The spatial attention map is further multiplied with the amodal feature to obtain the attention amodal feature. This feature is then fed into a mask prediction module, which is structured with a 2 × 2 transposed convolutional layer employing a stride of 2, coupled with a 1 × 1 convolutional layer using a stride of 1, to derive the amodal mask M i aEmploying AIS protocols, the training adopts a two-stage instance segmentation process similar to Mask R-CNN, facilitating concurrent training of both bounding box and amodal mask prediction heads alongside the object detection framework. In essence, the training procedure optimizes a multi-task loss function L as follows:EQUATIONwhere L det is object detection loss, defined similarly to that in Faster R-CNN object detection. The occluding mask loss L o , the visible mask loss L v , the amodal mask loss L a , and the classification loss L cls are computed using cross entropy loss with the corresponding ground truth. 
datasets, metrics and implementation details: 
 Datasets: We benchmark our AISDiff on three AIS datasets, namely KINS [19] , COCOA-cls [3] , and D2SA Metrics: Following existing AIS methods [24, 26] , we adopt mean average precision (AP) and mean average recall (AR). Tab. 1 depicts the comparison between AISDiff and SOTA AIS methods on the KINS dataset. AISDiff demonstrates consistent improvements across various backbones, including ResNet [5] , ResNet [5] , and RegNet [21] . Specifically, when compared to methods utilizing ResNet-50 as the backbone, our method outperforms both SOTA methods that use prior (e.g., and VRSP-Net [26] by 1.8 AP) and methods that do not use shape prior (e.g., AIS-Former [24] by 0.1 AP), respectively. When ResNet-101 is utilized as the backbone, our method achieves a improvement over AISFormer, outperforming it by 0.5 AP. Furthermore, compared to APSNet [14] and AISFormer [24] on the RegNet backbone, our approach achieves SOTA performance by surpassing them 0.5 AP. D2SA. Tab. 2 further validates our approach on D2SA dataset. We achieve best results across all metrics. Specifically, we gains 0.74 on AP and 0.12 AR in comparison with the second best method, i.e. VRSP-Net. COCOA-cls. Tab. 3 shows our results on COCOA-cls dataset. Our AISDiff also outperform other methods on all metrics. In fact, it outperforms the second best by 0.16 AP and 0.03 AR. put ROIs, Visible Masks, Occluding Masks, and Amodal Masks. Fig. 5 visualizes the spatial attention map of the Shape Prior Amodal Predictor on ROIs of the image. The attention maps are well-constrained to the object shape. Moreover, we can see that the decoder typically attends to the visible parts of objects that are similar to the occluded regions when predicting the amodal mask. 
conclusion: 
 In conclusion, we propose AISDiff, an AIS mask head with a Diffusion Shape Prior Estimation module. This module, termed DiffSP, leverages pre-trained conditioned diffusion models on extensive datasets to extract nuanced visual features for deriving the shape prior of the object. Furthermore, we present the Shape Prior Amodal Predictor, which utilizes attention-based feature maps from the shape prior to enhance amodal segmentation. Through extensive experimentation across diverse AIS benchmarks, we affirm the efficacy of AISDiff. 
 
 Appendix:  


Decision: This paper introduces a novel approach using diffusion models for amodal instance segmentation (AIS), showing promising results on multiple benchmarks. It enhances AIS performance by leveraging large-scale pre-trained models for shape estimation. While the method is an improvement over existing techniques, the marginal performance gain and lack of comparison with recent relevant works are noted weaknesses. Also, the relationship between humans and segmentation in city scenes should be discussed to better match this workshop's topic.
Overall, the approach is innovative and the paper is well-written, warranting acceptance with some reservations regarding originality and computational efficiency analysis.

Reviews: review: title: title: Review of Amodal Instance Segmentation with Diffusion Shape Prior Estimation
main_review: main_review: ** Summary **

This paper proposes a method to address Amodal Instance Segmentation. To mitigate the overfitting issue of relying on shape prior information from previous research, this approach uses a large-scale pretrained diffusion model. By introducing a new architecture with the Diffusion Shape Prior Estimation module, it enhances the performance of amodal instance segmentation.

** Pros **

- (Significance) The paper addresses a significant problem of amodal instance segmentation (AIS) that AIS has various applications.
- (Clarity)  The introduction of prior work and the clear presentation of methods and experiments make it easy to read.
- (Method) The authors effectively leverage the powerful ability of large-scale pretrained models for shape prior estimation, presenting a logically sound approach that is easy to follow. 
- (Effectiveness) Performance has been proven across various benchmarks, including KINS, COCOA-cls , and D2SA.
 							
** Cons **

- (Originality) While the novelty is not surprisingly remarkable, the addition of a new module based on existing work has improved performance. 
- (Effectiveness) Considering the complex added module DiffSP and Shape Prior Amodal Predictor, the performance improvement is somewhat marginal compared with the methods [14, 24] which do not require shape prior.
- (Minor correction): Line283: four AIS datasets -> three AIS datasets.
title: title: Official Review
main_review: main_review: * Summary

This paper addresses amodal instance segmentation (AIS), which aims to recover wholes from occluded parts in images. AIS methods typically leverage shape priors for segmentation. However, prior works often rely on shape priors learned from a limited number of unlabeled masks, which are prone to overfitting and neglect class information. To address this, the paper proposes using text-to-image (T2I) diffusion models trained on a large-scale dataset with semantics. The proposed method infers class and condition it into the T2I diffusion to obtain shape priors, which are used to generate final AIS masks. This approach outperforms prior works not using diffusion priors.


* Strength

The 3D or shape prior abilities of T2I diffusion models have been reported in many research papers. Therefore, it is reasonable to explore how to properly utilize the information for the AIS problem.


* Weakness

The baseline methods are outdated; all of them are from before 2022. There are a bunch of works that utilize diffusion priors for the AIS problem [1, 2, 3]. The paper should compare with those approaches both methodologically and empirically.

[1] Ozguroglu et al. pix2gestalt: Amodal Segmentation by Synthesizing Wholes. CVPR 2024.\
[2] Zhan et al. Amodal Ground Truth and Completion in the Wild. CVPR 2024.\
[3] Xu et al. Amodal Completion via Progressive Mixed Context Diffusion. CVPR 2024.
title: title: Review
main_review: main_review: Summary:

This paper introduces Diffusion-based methods into Amodal Instance segmentation for the first time. By leveraging the shape knowledge, the authors propose Diffusion Shape Prior Estimation (DiffSP) to estimate the potential shape knowledge of various object for amodal mask prediction. Moreover, the proposed method can be applied on various detectors and show the effectiveness on three different datasets. 


Strength:

1, The first work that explores the diffusion model for amodal segmentation and mask prediction.

2, The motivation is clear since diffusion models are pre-trained with huge object level data which can be used to predict missing masks or shape of amodal objects.

3, The overall writing is good and easy to understand.

4, The proposed method achieves consistent improvements on various detectors and datasets, which is convincing to me.


Weakness:

1, Missing analysis on computation cost or run time analysis. 

2, Better to show different diffusion model’s effects, such as Stable-XL.