SmartBrush: Text and Shape Guided Object Inpainting with Diffusion Model

Shaoan Xie,Zhifei Zhang,Zhe Lin,Tobias Hinz,Kun Zhang

DOI: https://doi.org/10.48550/arXiv.2212.05034

2022-12-10

Abstract:Generic image inpainting aims to complete a corrupted image by borrowing surrounding information, which barely generates novel content. By contrast, multi-modal inpainting provides more flexible and useful controls on the inpainted content, \eg, a text prompt can be used to describe an object with richer attributes, and a mask can be used to constrain the shape of the inpainted object rather than being only considered as a missing area. We propose a new diffusion-based model named SmartBrush for completing a missing region with an object using both text and shape-guidance. While previous work such as DALLE-2 and Stable Diffusion can do text-guided inapinting they do not support shape guidance and tend to modify background texture surrounding the generated object. Our model incorporates both text and shape guidance with precision control. To preserve the background better, we propose a novel training and sampling strategy by augmenting the diffusion U-net with object-mask prediction. Lastly, we introduce a multi-task training strategy by jointly training inpainting with text-to-image generation to leverage more training data. We conduct extensive experiments showing that our model outperforms all baselines in terms of visual quality, mask controllability, and background preservation.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

### Problems the paper attempts to solve

This paper aims to solve three main challenges in image inpainting:

1. **Mask Misalignment**:
- When using shape - guided methods in existing multi - modal image inpainting, the generated objects may not fully follow the given mask shape, especially when using coarse masks.

2. **Text Misalignment**:
- When using text - guided methods in existing approaches, the generated content may be inconsistent with the text description, especially when there is a semantic mismatch between local content and the global text description.

3. **Background Preservation**:
- Existing methods often lead to distortion or change in the background area during the inpainting process, especially when using coarse masks.

To solve these problems, the authors propose a new diffusion model **SmartBrush**, which combines text and shape guidance and introduces a precision - control mechanism. Specifically, the model solves the above problems in the following ways:

- **Text and Shape Guidance**:
- The model uses not only the text description to guide the generation of objects but also the shape mask to constrain the shape of the generated objects. This makes the generated content more in line with users' expectations.

- **Mask Precision Control**:
- By generating masks of different precisions (from fine to coarse), the model allows users to specify how closely the generated objects should follow the given mask shape. This provides users with more flexibility.

- **Background Preservation**:
- The model is encouraged to predict an accurate foreground object mask during the training process, so as to better preserve background information during the inpainting process. This ensures that the generated objects will not cause unnecessary changes to the surrounding background.

### Main contributions

1. **Introduced a diffusion model for object inpainting with text and shape guidance**:
- The model can perform inpainting according to object masks of different precisions, achieving a new level of control.

2. **Preserved the image background by predicting the foreground object mask**:
- Especially when using coarse input masks, the model can better preserve the original background.

3. **Trained with instance segmentation masks and local text descriptions**:
- Compared with random masks and texts describing the entire image, this method can better align local content and text descriptions.

4. **Proposed a multi - task training strategy**:
- By jointly training object inpainting and text - to - image generation tasks, more training data are utilized and the performance of the model is improved.

### Experimental results

The authors conducted extensive experiments on multiple datasets. The results show that **SmartBrush** outperforms existing baseline methods in terms of visual quality, text alignment, and background preservation. User studies also indicate that users are more likely to choose the generation results of **SmartBrush**.