{'fraction_non_alphanumeric': 0.025392457314403653, 'fraction_numerical': 0.01244605666395647, 'mean_word_length': 4.72093023255814, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Segmentation in Style: Unsupervised Semantic Image Segmentation with Stylegan and CLIP

We introduce a method that allows to automatically segment images into semantically meaningful regions without human supervision. Derived regions are consistent across different images and coincide with human-defined semantic classes on some datasets. In cases where semantic regions might be hard for human to define and consistently label, our method is still able to find meaningful and consistent semantic classes. In our work, we use pretrained StyleGAN2 generative model: clustering in the feature space of the generative model allows to discover semantic classes. Once classes are discovered, a synthetic dataset with generated images and corresponding segmentation masks can be created. After that a segmentation model is trained on the synthetic dataset and is able to generalize to real images. Additionally, by using CLIP we are able to use prompts defined in a natural language to discover some desired semantic classes. We test our method on publicly available datasets and show state-of-the-art results.


Introduction
The development of deep convolutional neural networks [5,6,18] has fueled remarkable progress in semantic segmentation and pushed state-of-the-art on a variety of datasets [2,19,21]. While the introduced methods show impressive results, their applicability is limited by a very slow process of image annotation [4]. This issue is more pronounced in the medical imaging community where a domain knowledge is required to perform annotations. Consistent labeling can be a difficult task for a team of annotators 1 https://github.com/warmspringwinds/ segmentation_in_style working on the same dataset as most semantic classes lack clear visual boundaries. This can lead to some heterogeneity in the labels which can affect the overall model's performance. At the same time most deep learning approaches require a considerable amount of training data to reach their best performance [5]. In this paper, we introduce a method that does not require any labeling, allows to discover consistent semantic regions across images, outperforms semisupervised methods that require a small number of annotations and is competitive with fully supervised image segmentation methods.
In our work we compare our method against semisupervised learning approaches (SSL) which use a small number of annotations. This serves as a very strong baseline for our approach and we show that we outperform recent state-of-the-art methods [10,20] even though we do not use any annotations. Our method allows to suggest consistent semantic classes and we show that on some datasets they coincide with annotations provided by human which allows us to measure its accuracy. As it can be seen our method allows to discover high quality semantic regions that coincide with human defined regions (see Fig. 1) and also suggest consistent semantic regions that might be hard for human to discover and label (see Fig. 2). Additionally, we introduce a method that allows us to use natural language to discover certain rare semantic regions like beard, glasses and hats (see Fig. 6). Finally, we show how we can automatically classify each of the discovered semantic regions and assign a prompt defined in a natural language to each of them (see Fig. 5). To demonstrate applicability of our method across other domains, we carry out experiments on an eye segmentation dataset [13] and show that our discovered semantic regions replicate human annotations (see Fig. 8) and that the accuracy of our provided segmentation masks is com-petitive with fully supervised methods.

Method
Our approach is based on StyleGAN2 [8] a very powerful generative model capable of generating realistic samples of objects of interest once it has been trained on a related dataset (Section 2.1). In our work we either use a pretrained models in case of human faces, cats, dogs and cartoons or train a model from scratch on a desired dataset in case of eye segmentation problem [13].
We show that by clustering in the feature space of trained StyleGAN2 model, we can discover meaningful semantic regions that are consistent across different samples of the model (Section 2.2).
To discover rare semantic classes that were not found during clustering since they appear relatively rarely in samples of StyleGAN2, we apply recently introduced approach for image manipulation [14,16] using text prompts with CLIP [15]. By manipulating latent vectors of sampled images, we make a desired semantic region to appear in almost every sample, which allows it to appear as one of the clusters (Section 2.3).
To conclude the process and attribute each semantic region with a prompt defined in a natural language, we modify the original CLIP [15] image encoder to accept image regions instead of a whole image. Then we encode each of the discovered semantic regions and classify them (Section 2.4).
Once a StyleGAN2 model is modified to deliver samples along with segmentation masks, we apply a simple knowledge distillation procedure by training a segmentation network on synthetic images and segmentation masks and show that it generalizes to real images (Section 2.5).

StyleGAN2
StyleGAN2 is used in our method as a generative backbone as it is able to generate images with high quality, gives us access to feature maps of generated images that we use for clustering and allows generated images to be manipulated which we use to discover additional semantic classes. Sampling is performed by drawing from a normal distribution z ∈ Z. A separate mapping network maps z to an intermediate latent vector w ∈ W . w is then transformed into k vectors that are used by different layers of generator network of StyleGAN2 (see Fig. 3). StyleGAN2 gradually generates feature maps of higher resolution as they get upsampled and processed by successive layers of the generative model.

Clustering
During clustering, we generate N samples and save the generated images as well as corresponding feature maps from a certain layers of StyleGAN. We mostly use feature maps from the seventh or ninth layer of the generative model. Using earlier layers of StyleGan results in more coarse but more semantically meaningful clusters. Later layers result in more fine clusters with less semantic meaning. We also performed experiments with concatenated features from different layers but did not notice any improvement. Feature maps have dimensionality of either N × 64 × 64 × 512 or N × 128 × 128 × 256 where the first number represents a number of samples the second and third numbers represent the spatial resolution and the last number represents the dimensionality of the feature space. We flatten the acquired feature maps and end up with N * 64 * 64 feature maps of size 512 or with N * 128 * 128 feature maps of size 256. After that we perform clustering by minimizing the following loss function: Where x x x n represents a single data point of size 512 or 256, r nk ∈ 0, 1 is a binary indicator variable, where k = 1, . . . , K describes which of the K clusters the data point x x x n is assigned to.N is equal to N * 128 * 128 or N * 64 * 64 depending on the layer of choice. We search for set of values {r nk } and {µ µ µ k } that minimize J. To solve this problem we use K-means clustering algorithm [12]. Discovered classes can be seen in Fig. 2. After the clusters are discovered we notice that new samples have the same semantic regions if we assign their features to the closest cluster. This gives us a way to generate image and corresponding segmentation masks at the same time with minimal additional computational overhead.

Class Discovery with Image Manipulation
Some classes that are rarely present in samples of Style-GAN like hat, beard or glasses are not always discovered by our method as is. In order to make desired semantic classes more present in samples that we use for clustering, we update a latent code of every image so that it contains a desired attribute. Let w ∈ W denote a latent code of a generated sample, and F(w) the corresponding generated image (where F is a generator network of StyleGAN). We employ a recently introduced method [14,16]: we collect a number of latent vectors and corresponding generated images and classify generated images with CLIP [15] given a natural language prompt like "a person with a beard". Given a set of pairs {(w i , b i )}, where w is a latent vector and b is a binary variable indicating if the generated images contains desired attribute, we learn a manipulation direction G, such that F (w + αG) yields an image where that attribute is introduced or amplified. The manipulation strength is controlled by α. The method depiction can be seen in Fig. 4.  Clustering allows us to find semantic regions of the generated images. After the clusters are found, a much bigger set of images is generated and using their feature maps we attribute each pixel to one of the previously discovered semantic clusters. A segmentation network is later on trained on the synthetic dataset.   Figure demonstrates a process of creating a synthetic annotations for rare classes that were not discovered during clustering. First, a latent direction in the Stylegan is learnt that adds a desired semantic class to almost every generated sample using the method of [14,16]. In this case, the vector G represents the a text promt "a person with glasses". After that, almost every sample has the desired attribute and it naturally appears as one of the clusters. As it can be seen, semantic region representing glasses is indeed represented by one of the clusters.
An example of latent code manipulation along the beard direction can be seen in Fig. 9 and discovered corresponding semantic regions are depicted in Fig. 6.

Cluster Classification
Once all the desired clusters are discovered, we need to assign a semantic class to each of them. In order to do so, we use pretrained text encoder and image encoder of the CLIP model [15]. We use text encoder as is and encode desired classes given corresponding text prompts like "hair", "forehead" etc (see Fig. 5). As for the image encoder we use a model based on residual neural network [5,15]. Our choice is motivated by the fact that this model can be modified to deliver embedding with bigger spatial resolution by Given a set of text prompts of desired classes and a set of generated images with corresponding clusters, we embed both text and image regions using pretrained text encoder and image encoders of CLIP [15]. After that we compute pairwise dot products between text and cluster embeddings. Each cluster is assigned to a text prompt that results in a biggest dot product value. For example, in a given set of images all clusters containing hair will be classified as "hair". removing downsampling layers and adding dilation factors to certain convolutional layers [2,19]. We observed that bigger spatial resolution gives better results for smaller clusters. Once we get embeddings for each pixel of every image, we average them per each cluster over many images. Averaging over many images also gives better results according to our experiments. Finally, we compute the dot products between embedding of each cluster and each text prompt: we assign the cluster to a class that resulted in a biggest dot product value. This way we use original clusters as region proposals that we classify.

Knowledge Distillation
Once we are able to generate synthetic images and corresponding annotations, we need a way to get the same segmentation masks for real images. Inspired by recent work in knowledge distillation for image manipulation with Style-GAN [17] we generate a synthetic dataset with images and corresponding segmentation masks and train a simple segmentation model on that dataset. We apply a segmentation method based on resnet-18 and dilated convolutions [19] and show that trained model also generalizes to real images as it can be seen in Fig. 7

Datasets and Results
For faces, we evaluate our model on CelebA-Mask8 dataset [11], which contains 8 part categories and compare our method to state-of-the-art semi-supervised methods [10,20]. As can be seen in the Table 1 our method achieves state-of-the-art results. In order to evaluate our method on face images that contain a bigger set of labels we test our method on CelebAMask-HQ dataset [9] segmentation face parsing dataset that contains 19 semantic classes. To simplify the dataset we merge left and right eye classes into one and do similar simplification for ears. We also eliminate earring and necklace attributes as StyleGAN struggles to generate these attributes realistically which results in poor performance. The results compared to fully supervised method can be found in Table 2. As it can be seen our method is performing worse than the fully supervised method but delivers overall good results. We also collect a small dataset where we label beard in order to test the accuracy of our method on this discovered semantic class. The results compared to fully supervised method can be found in Table 4. To test if our method is able to generalize to completely different domains, we use our method to perform eye segmentation on OpenEDS [13] dataset. We train Style-GAN2 on images of the dataset from scratch and apply our algorithm to the trained model. As it can be seen, overall iris, pupil and sclera were discovered which coincides with Figure 6. Example of synthetic image and annotation pairs created for rare classes with our method with CLIP [15] using text prompts "a person with a beard" (first row), "a person with glasses" (second row), "a person wearing a hat" (third row). Interestingly, the discovered class representing glasses does not include eye regions which is consistent with human annotations of glasses for some datasets like CelebAMask-HQ [9]. semantic classes of the original dataset (see Fig. 8). Moreover, our algorithm achieves a segmentation performance which is very close to the fully supervised method in this case, since the semantic classes are simpler (see Table 3).

Conclusion
We propose a powerful method for unsupervised learning that allows to discover meaningful and consistent semantic classes which mostly coincide with classes defined and labeled by human. Our method is successful in cases where it might be hard for human to define and consistently label semantic classes. We propose a way to derive semantic class annotations by defining a text prompt which allowed us to discover classes like beard that currently has no public annotations. We show that training a segmentation model on our generated synthetic images along with segmentation masks generalizes to real images and shows competitive results with fully supervised methods. Our method achieves state-of-the-art results as it performs better than recently introduced semi-supervised models.

Method
DatasetGAN [20] Table 2. Results of our algorithm on CelebAMask-HQ [9] in comparison to fully supervised method. The dataset has additional face semantic classes compared to CelebA-Mask8 Dataset [11].

Method
SegNet [13] Ours Number of Manually Annotated Images Used All 0 Mean IOU 84.1 82.39 Table 3. Results of our algorithm on OpenEDS 2020 eye segmentation dataset [13] in comparison to fully supervised method. The dataset has 4 semantic classes. As it can be seen, our method has a very similar performance compared to fully supervised method as semantic classes are simpler in this dataset.

Method DRN [13] Ours Number of Manually Annotated Images Used
All 0 Mean IOU 88.8 84.29 Table 4. Results of our algorithm on our custom beard segmentation dataset compared to fully supervised approach. Figure 8. Example of semantic classes discovered by clustering features from different layers for OpenEDS eye segmentation dataset. As it can be seen, overall iris, pupil and sclera were discovered which coincides with semantic classes of the original dataset.