Where's Waldo: Diffusion Features For Personalized Segmentation and Retrieval

Dvir Samuel1,2, Rami Ben-Ari2, Matan Levy3, Nir Darshan2, Gal Chechik1,4,
1Bar Ilan University, 2OriginAI, 3The Hebrew University of Jerusalem, 4NVIDIA Research
NeurIPS 2024
Paper Code & Data

Personalized segmentation task involves segmenting a specific reference object in a new scene. Our method is capable to accurately identify the specific reference instance in the target image, even when other objects from the same class are present. While other methods capture visually or semantically similar objects, our method can successfully extract the identical instance, by using a new personalized feature map and fusing semantic and appearance cues. Red and green indicate incorrect and correct segmentations respectively.


Abstract

Personalized retrieval and segmentation aim to locate specific instances within a dataset based on an input image and a short description of the reference instance. While supervised methods are effective, they require extensive labeled data for training. Recently, self-supervised foundation models have been introduced to these tasks showing comparable results to supervised methods. However, a significant flaw in these models is evident: they struggle to locate a desired instance when other instances within the same class are presented. In this paper, we explore text-to-image diffusion models for these tasks. Specifically, we propose a novel approach called PDM for Personalized Diffusion Features Matching, that leverages intermediate features of pre-trained text-to-image models for personalization tasks without any additional training. PDM demonstrates superior performance on popular retrieval and segmentation benchmarks, outperforming even supervised methods. We also highlight notable shortcomings in current instance and segmentation datasets and propose new benchmarks for these tasks.


Are instance features even encoded in a pre-trained text-to-image model?

(a) Apperance Features: We found that instance appearance features are encoded in the queries (\(\mathcal{Q}^{SA}\)) and keys (\(\mathcal{K}^{SA}\)) matrices of the self-attention block. Here we show PCA visualization of features obtained from the first self-attention block in the last layer of the U-Net module, at various diffusion timesteps. Objects with similar textures and colors have similar features. The dog's color in \(I_1\) is similar to the colors of both the dog and the cat in \(I_2\), indicating textural similarity. Additionally, the localization is sharper at larger timesteps.

(b) Semantic Features: Visualization of the cross-attention map for a given prompt "dog". Note the higher region correlation (brighter colors) corresponding to the dog, while overlooking the cat in the bottom image.


PDM: Personalized Diffusion Features Matching

An overview of our Personalized Diffusion Features Matching approach. PDM combines semantic and appearance features for zero-shot personalized retrieval and segmentation. We first extract features from the reference, \(I_r\) and target \(I_t\) images. Appearance similarity is determined by dot product of cropped foreground features from the reference feature map, \(F_r^{AM}\) and the target feature map \(F^A_t\). Semantic similarity is calculated as the product between class name token \(C\) and the target semantic feature map \(F^S_t\) to create a Semantic Map. The final similarity map combines both maps by average pooling.


Personalized Segmentation

Red and green indicate incorrect and correct segmentation, respectively. Our method accurately recognizes the reference instance despite significant variations (view angle, pose, or scale), while other methods often capture false positives from the same category.



Personalized Retrieval

Top-1 retrieved image is shown for each method. Note how our model identifies images containing the same instance, despite their small size and large variations. Other methods tend to capture only semantic similarity.

BibTeX

@article{Samuel2024Waldo,
  title={Where's Waldo: Diffusion Features For Personalized Segmentation and Retrieval},
  author={Dvir Samuel and Rami Ben-Ari and Matan Levy and Nir Darshan and Gal Chechik},
  journal={NeurIPS},
  year={2024}
}