[Paper Review] Learning Vector Quantized Shape Code for Amodal Blastomere Instance Segmentation

Won-Dong Jang, Donglai Wei+7 more|arXiv (Cornell University)|Dec 2, 2020

Reproductive Biology and Fertility4 citations

TL;DR

This paper proposes a novel amodal instance segmentation method for blastomeres using a Vector Quantized Variational Autoencoder (VQ-VAE) to learn discrete shape codes from ground-truth amodal masks, improving robustness to occlusion. By classifying region features into these shape codes and refining them with a dedicated decoder, the method achieves superior segmentation accuracy on both embryo and natural image benchmarks, outperforming state-of-the-art approaches in mAP and IoU metrics.

ABSTRACT

Blastomere instance segmentation is important for analyzing embryos' abnormality. To measure the accurate shapes and sizes of blastomeres, their amodal segmentation is necessary. Amodal instance segmentation aims to recover the complete silhouette of an object even when the object is not fully visible. For each detected object, previous methods directly regress the target mask from input features. However, images of an object under different amounts of occlusion should have the same amodal mask output, which makes it harder to train the regression model. To alleviate the problem, we propose to classify input features into intermediate shape codes and recover complete object shapes from them. First, we pre-train the Vector Quantized Variational Autoencoder (VQ-VAE) model to learn these discrete shape codes from ground truth amodal masks. Then, we incorporate the VQ-VAE model into the amodal instance segmentation pipeline with an additional refinement module. We also detect an occlusion map to integrate occlusion information with a backbone feature. As such, our network faithfully detects bounding boxes of amodal objects. On an internal embryo cell image benchmark, the proposed method outperforms previous state-of-the-art methods. To show generalizability, we show segmentation results on the public KINS natural image benchmark. To examine the learned shape codes and model design choices, we perform ablation studies on a synthetic dataset of simple overlaid shapes. Our method would enable accurate measurement of blastomeres in in vitro fertilization (IVF) clinics, which potentially can increase IVF success rate.

Motivation & Objective

• To address the challenge of accurate amodal instance segmentation in time-lapse embryo imaging for in-vitro fertilization (IVF) applications.
• To improve robustness to occlusion by learning a discrete shape prior via VQ-VAE, rather than direct pixel-level mask regression.
• To enhance bounding box quality and segmentation fidelity by integrating an occlusion map prediction branch.
• To demonstrate generalizability of the method beyond embryonic cells to diverse object classes in natural scenes.
• To enable automated, accurate measurement of blastomere shape and symmetry, potentially increasing IVF success rates.

Proposed method

• Pre-train a VQ-VAE on ground-truth amodal masks to learn a discrete embedding space representing complete object shapes.
• Integrate the VQ-VAE into a detection and segmentation pipeline, replacing direct mask regression with classification into learned shape codes.
• Use a refinement decoder to reconstruct high-resolution amodal masks from the discrete shape codes and low-level features.
• Incorporate an occlusion map prediction head to guide the model with visibility information, improving detection under partial occlusion.
• Train the full network with a multi-task loss combining embedding loss, occlusion loss, and mask reconstruction loss.
• Apply the model to both embryo-specific and general-purpose benchmarks (e.g., KINS) to validate performance and generalization.

Experimental results

Research questions

• RQ1 Can a discrete shape code learned via VQ-VAE improve amodal instance segmentation robustness under varying occlusion levels?
• RQ2 Does classifying features into shape codes outperform direct regression of amodal masks in terms of segmentation accuracy?
• RQ3 How effective is the integration of an occlusion map branch in improving bounding box and mask prediction for partially visible objects?
• RQ4 To what extent does the proposed method generalize to natural image datasets with diverse object shapes?
• RQ5 What is the contribution of each component (VQ-VAE, refinement decoder, occlusion map) to overall performance?

Key findings

• On the internal embryo benchmark, the proposed method outperforms previous state-of-the-art methods in amodal segmentation accuracy.
• On the KINS natural image benchmark, the method achieves a mAP of 0.303 with Mask R-CNN, surpassing the baseline FCN (0.303 vs. 0.293) and matching Mask R-CNN + ASN (0.315).
• The ablation study confirms that all components—VQ-VAE, refinement decoder, and occlusion map—are essential, as removing any reduces mAP from 0.303 to 0.298 and 0.281.
• In the occlusion robustness test, the proposed method achieves an IoU of 0.968, significantly outperforming FCN (0.935), with sharper and more complete shape completions.
• Qualitative results show that the proposed method reliably reconstructs full object shapes even under heavy occlusion, while baseline FCN produces blurry, soft predictions on unseen regions.
• The method successfully segments occluded cars in KINS, though failure cases occur due to non-maximum suppression removing overlapping bounding boxes.