S. Sun, M. Xian, F. Xu, T. Yao, L. Capriotti
{Cascade-Forward Refinement (CFR); Iterative Click Loss (ICL); SUEM Image Augmentation}
Paper: https://arxiv.org/pdf/2303.05620v1.pdf
Code: https://github.com/TitorX/CFR-ICL-Interactive-Segmentation
Interactive image segmentation
Extracts desired objects from images using user clicks.
Critical for real-world application of deep learning models.
Reduces manual labeling costs.
Challenge
Inflexible inference: Rigid rules and separate refinement models limit adaptability.
Training imbalance: Difficulty balancing user clicks needed for good performance with training the model effectively.
Propose a click-based and mask-guided interactive image segmentation framework.
Three novel components:
Cascade-Forward Refinement (CFR): A unified inference framework to generate segmentation results in a coarse-to-fine manner.
Iterative Click Loss (ICL): A unique loss function that simulates user clicks during training. ICL trains the model to not only improve segmentation quality but also decrease the need for multiple user clicks by focusing on uncertain regions.
SUEM Image Augmentation: Effective data augmentation strategies for interactive segmentation. SUEM generates a range of semi-synthetic masks which help the model understand different user clicking styles.
GrabCut [20].
NC-Cut [25].
EISeg [24]: Connectedness-based approach.
DIOS [27] encoded the background and foreground user clicks into two distance maps and concatenated the image with the two maps as input.
BRS [10] and f-BRS [22] formulated the task as an inference-time online optimization problem.
BRS optimizes the image or distance maps during inference to improve the segmentation results for given user clicks.
f-BRS optimizes the intermediate layer weights to achieve the same goal and speed up the inference.
FCA-Net [14] built a first-click attention module that enhances the importance of the first click as it usually plays a critical role to identify the location of the main body.
UCP-Net [4] proposed a novel contour-based approach that asks the user to provide clicks on the contour of objects.
[17] ITIS An iterative sampling strategy samples a single point from the center of a misclassified area iteratively.
RITM [23] uses the previous output as model input to achieve higher segmentation quality.
FocalClick [2] proposes a coarse-to-fine pipeline with a SegFormer [26] backbone. (determine local regions that possibly contain errors and uses a local refinement network to update these regions).
SimpleClick [15] adopts the Plain Vision Transformer (Plain ViT), which was pretrained with MAE [9], as the backbone of the RITM approach.
==> Our work builds upon the RITM and SimpleClick approaches and further explores the nature of interactive segmentation.
[6] a simple copy-paste strategy can serve as a powerful data augmentation for instancelevel segmentation tasks.
==> We develop more specific image augmentation approaches for adapting interactive segmentation tasks.
New iterative mask-guided framework for interactive segmentation.
Features:
Inference-time refinement without extra modules.
Click-encoded training loss for fewer clicks.
Effective image augmentation for interactive scenarios.
Goals:
Efficiency.
Minimize user effort.
Optimize performance.
Input: Image, Click Maps, Previous Mask.
Interactive Segmentation Model:
Patch Embed.
ViT Blocks.
Feature Pyramid.
Segmentation Head.
Segmentation Result.
Cascade-Forward Refinement.
Cascade-Forward Refinement Loop.
User Interaction Loop.
SUEM Copy-paste.
Pre-training: SimpleClick (ViT-Base and ViT-Huge models), that utilize the Plain Vision Transformer [3] pretrained with MAE.
Fine-tuning the pre-trained models using our proposed methods.
Segmentation Head.
Iterative Click Loss (ICL).
Normalized Focal Loss (NFL) is used with α = 0.5 and γ = 2.
Iterative Click Learning (ICL)
3 clicks are generated iteratively for training the model.
βi ∈ [1, 2, 3].
Optimizer and Learning Rate:
Adam optimizer is used with β1 = 0.9 and β2 = 0.999.
Learning rate is set to 5 × 10^−6.
Fine-tuning:
All models undergo 1 epoch of fine-tuning from the SimpleClick pre-trained models.
Batch sizes: 140 for ViT-Base, 32 for ViT-Huge.
The core idea behind CFR is to streamline the inference process of interactive image segmentation by generating progressively refined segmentation predictions in a coarse-to-fine manner.
Develop a single-model inference strategy for iterative segmentation refinement, having two inference loops:
Outer loop - Initial stage (Coarse): Generates coarse segmentation masks using incremental user interactions.
Input: Original image and a few user clicks.
Output: A coarse segmentation mask that roughly outlines the object of interest.
Inner loop - Refinement stage (Fine): Refines the masks by forwarding the segmentation model multiple times with the same input image and user clicks.
Input: Coarse mask, original image, user clicks.
Output: Refined mask.
Let:
f denotes a deep neural network for segmentation.
X is the input image.
Yt is the output of the t-th step (mask).
The sequence of user clicks (maps) at the t-th step is denoted as:
where (uk, vk) represents the coordinates of a user click and lk ∈ {0, 1} denotes its label (0 for background and 1 for foreground).
The Cascade-Forward Refinement approach iteratively updates the coarse segmentation masks and refinement masks using:
where Y0t denotes the output of the t-th coarse segmentation step (outer loop), and Y0t-1 is the last refined mask at step t-1.
Eq. 2 defines the refinement results at the i-th step in the inner loop, and n defines the number of refinement steps which could be defined as a fixed number or determined adaptively.
No additional user clicks are required during the refinement process (inner loop), and the segmentation mask is continuously refined to provide higher-quality input.
A fixed-step CFR applies n times refinement for each user click. Improves output quality without model modification.
Adaptive CFR (A-CFR) counts the number of altered pixels between Yit and Yi-1t and terminates the inner loop when the number of changed pixels falls below a specified threshold, or the maximum step is reached.
Recent works:
[22] encode all generated user clicks for each image into two disk maps and input them into a segmentation network during training.
[17, 23] adopted a hybrid strategy that combined randomly sampled clicks with iteratively generated clicks to generate input maps.
Problem:
The trained models may need more user interactions during inference
Have no effective way to balance the number of inputted clicks and segmentation performance.
Iterative Click Loss (ICL) approach embeds the number of clicks during training.
An initial set of randomly sampled clicks, denoted as P0, is generated from the ground truth to forward the model and obtain an initial output Y0.
Clicks are generated iteratively (one click at a time) by sampling from the misclassified regions in Y0, and the newly produced click is combined with P0 to form a new sequence P1.
This process is repeated to generate the whole click sequence Pt.
Formular of click-sampling strategy:
where 0 is the zero-initialized mask, Y is the ground truth, and S is a sampling function [23] that proposes a single click among the misclassified areas in the output mask.
A conventional total loss function [23]:
where L is the Normalized Focal Loss [21].
Note that in Eq. 4, a sequence of click sets [P0, P1, ..., Pt] and a sequence of outputs [Y0, Y1, ..., Yt-1] are used to generate segmentation mask Yt; however, only the final output is applied to calculate the Normalized Focal Loss and update the model parameters.
In the proposed ICL approach, a new loss function is built to accumulate the weighted losses of the generated mask sequence.
where βi is used to control the weight of each term.
Each click produces one loss term in the above equation, and minimizing the loss improves the segmentation performance and reduces the number of clicks simultaneously.
By increasing the weights of the loss term for more clicks (larger i), the model is incentivized to use fewer clicks to achieve accurate segmentation.
SUEM Copy-paste (C&P): to generate large and diverse datasets for interactive segmentation scenarios.
Principle: inserting randomly selected objects from one image into another.
Source object and Source image: the object of interest and its corresponding image, respectively.
The extra object and extra image: the object and its corresponding image selected randomly from the training set.
Four methods:
Simple C&P Mode. The simple copy-paste mode involves inserting a source object into a randomly selected extra image, with the mask of the source object serving as the ground truth.
Union C&P Mode. Interactive segmentation typically involves identifying a target object that comprises multiple objects, such as a man embracing a child. To simulate such a scenario, we employ the union copy-paste mode, where an extra object is pasted into the source image, and the ground truth is determined as the union of their respective masks.
Exclusion C&P Mode. Another scenario that frequently arises is that the object of interest is obstructed by another object, such as a person standing behind a pole. To address this issue, we introduce the exclusion copy-paste mode, where an extra object is copied into the source image, and the mask of the source object is utilized as the ground truth, excluding the mask of the extra object.
Image Mixing Mode. The approach of image-mixing involves blending a source image with an extra image and utilizing the mask of the source object as the ground truth.
GrabCut (50 images): Simple dataset with one instance per image.
Berkeley (96 images): Small dataset with 100 instances.
DAVIS (345 images): High-quality segmentation masks extracted from videos.
Pascal VOC (1449 images): Validation set used for evaluation.
SBD (8497 train, 2857 val images): Larger dataset for training and validation.
COCO+LVIS (C+L, 104k images): Combined dataset used exclusively for training, offering extensive data for model learning.
Number of Clicks (NoC), measuring user inputs needed for satisfactory segmentation.
Thresholds: NoC@90 and NoC@95 (more stringent than commonly used NoC@85 and NoC@90) to ensure high-quality segmentation.
Click generation: Method from [12] used for evaluation.
Click limit: 20 clicks per image for efficient assessment.
The current implementation of ICL encodes the order and number of clicks indirectly during training. The order of user clicks embeds user intentions and perspectives which might be useful for guiding the interactive segmentation approaches.
The current disk map used to encode user clicks is inadequate for representing the order. A novel form of click encoding similar to language model encoding could be expected. For example, a sequential model could be applied to encode the order information of a sequence of clicks into a map.
