Object Detection

Once we get the object class to be replaced by parsing the user prompt, we use a MaskRCNN model pretained on the COCO dataset to predict the bounding box of the specified object in the input image. The above figure shows the model structure of MaskRCNN. MaskRCNN is an extension of Faster R-CNN with an additional branch for predicting segmentation masks on each region of interest. The model has two stages. During the first stage, regions of interest are proposed by the region proposal network. These will be further processed by the Region of Interest Pooling networks to exact features and perform object classification. In parallel to object classification, binary masks of objects are predicted by the Mask Head portion denoted in the figure. The outputs of MaskRCNN include class labels, bounding boxes, and object masks for all objects. In OpenGIT, we make use of the bounding box of the specified object for later processing.

Class-Conditioned Image Editing

Given an input image, target label and a bounding box automatically produced by the previous MaskRCNN Pipeline, we use a pretrained MaskGIT model for class-conditioned image editing. The region inside of the bounding box is considered as the initial mask to the MaskGIT iterative decoding algorithm.

The iterative decoding algorithm runs T steps. It starts with the initial mask, in our case, it's the rectangle gray area inside of the bounding box. Let the masked tokens be $Y_M^{(t)}$. At each step t, the model produces the image via predicting all tokens at the same time but only preserving the most confident ones. We define each masked location as $i$. A token $y_i^{(t)}$ is sampled based on its prediction probabilities $p_i^{(t)} \in \mathbf{R}^K$. The confidence score is defined as the corresponding prediction score of a token $y_i^{(t)}$, and it indicates how confident our model is of this prediction. For the non-masked regions, the confidence score is set to 1.0. For the less confident tokens, they are considered as the mask input to the next iteration, the mask ratio is designed to be decreasing so that within T iterations, all tokens are being generated.

After we perform this decoding algorithm for T iterations, the model generates a resulting image that has contents of the given target label and contexts of the input image.

In our example above, we see that given our original input image of a dog, target label cat and a bounding box, the contents within a box becomes the initial mask, and we obtain a cat image with the contexts (e.g. fur and color) that's similar to the original dog image.

Fine-tuning with Text

Given a sample of class-conditioned image $x$, let $x = f_d(q_z)$ where $f_d$ is the VQ-VAE decoder and $q_z$ is the grid of quantized latent vectors. We optimize the text image similarity loss using pretrained CLIP model, CLIP($f_d(q_z)$, text_prompt). By fixing the decoder $f_d$, we optimize the text image similarity loss by performing gradient descent steps in $q_z$ space. This allows us to fine-tune the generated image on an open-domain text prompt. We see that pretrained CLIP model provides especially good guidance when we want to modify the color or texture of the generated images. However, if if we want to generate a new object from scratch, performing gradient descent on $q_z$ w.r.t. to CLIP loss does not lead to good results. In practice, we rely on pretrained MaskGIT to do content and object level image editing and we use CLIP optimization to control the texture of the output. When combined, MaskGIT + CLIP can produce a lot of creative edits in a huge space that is effectively 1,000 ImageNet classes * millions of potential texture and style adjective text prompts.

The Effect of Bounding Box Size

While generating edited images with OpenGIT, we find that the output is particularly sensitive to the size of the bounding box, even if the bounding box always bounds the same object. Therefore, we introduce another parameter, $r$, the ratio of the size of the desired bounding box to that of the original bounding box given by MaskRCNN as described above. While rescaling the bounding box, we make sure both the center and the aspect ratio of the original bounding box are preserved. Here are some results produced by OpenGIT when we specify different $r$ while keeping other inputs the same.

To generate the above visualizations, we give the same prompt "change cat to rabbit". For the first example, when $r = 0.4$, the bounding box only covers the center part of the cat's face. It's expected that the model would produce a rabbit face that merges well with the original image. When $r = 1$, the whole cat is covered in the bounding box. In this case, the model replaces the cat with a rabbit, including ears and the body. Since there is no inherent quality difference between the two ways of editing and it should be up to the user to decide which output gives the desired image editing effect, we let users specify $r$ directly in our interface. Future work includes inferring $r$ automatically given descriptive words like "face only", "partially", "completely", etc.