Spaces:
Runtime error
Runtime error
File size: 4,970 Bytes
82ef366 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
from typing import Tuple, List
import nltk
import numpy as np
import torch
from sklearn.cluster import KMeans
from constants import STYLE_INDEX, STRUCT_INDEX
nltk.download('punkt')
nltk.download('averaged_perceptron_tagger')
"""
Self-segmentation technique taken from Prompt Mixing: https://github.com/orpatashnik/local-prompt-mixing
"""
class Segmentor:
def __init__(self, prompt: str, object_nouns: List[str], num_segments: int = 5, res: int = 32):
self.prompt = prompt
self.num_segments = num_segments
self.resolution = res
self.object_nouns = object_nouns
tokenized_prompt = nltk.word_tokenize(prompt)
forbidden_words = [word.upper() for word in ["photo", "image", "picture"]]
self.nouns = [(i, word) for (i, (word, pos)) in enumerate(nltk.pos_tag(tokenized_prompt))
if pos[:2] == 'NN' and word.upper() not in forbidden_words]
def update_attention(self, attn, is_cross):
res = int(attn.shape[2] ** 0.5)
if is_cross:
if res == 16:
self.cross_attention_32 = attn
elif res == 32:
self.cross_attention_64 = attn
else:
if res == 32:
self.self_attention_32 = attn
elif res == 64:
self.self_attention_64 = attn
def __call__(self, *args, **kwargs):
clusters = self.cluster()
cluster2noun = self.cluster2noun(clusters)
return cluster2noun
def cluster(self, res: int = 32):
np.random.seed(1)
self_attn = self.self_attention_32 if res == 32 else self.self_attention_64
style_attn = self_attn[STYLE_INDEX].mean(dim=0).cpu().numpy()
style_kmeans = KMeans(n_clusters=self.num_segments, n_init=10).fit(style_attn)
style_clusters = style_kmeans.labels_.reshape(res, res)
struct_attn = self_attn[STRUCT_INDEX].mean(dim=0).cpu().numpy()
struct_kmeans = KMeans(n_clusters=self.num_segments, n_init=10).fit(struct_attn)
struct_clusters = struct_kmeans.labels_.reshape(res, res)
return style_clusters, struct_clusters
def cluster2noun(self, clusters, cross_attn, attn_index):
result = {}
res = int(cross_attn.shape[2] ** 0.5)
nouns_indices = [index for (index, word) in self.nouns]
cross_attn = cross_attn[attn_index].mean(dim=0).reshape(res, res, -1)
nouns_maps = cross_attn.cpu().numpy()[:, :, [i + 1 for i in nouns_indices]]
normalized_nouns_maps = np.zeros_like(nouns_maps).repeat(2, axis=0).repeat(2, axis=1)
for i in range(nouns_maps.shape[-1]):
curr_noun_map = nouns_maps[:, :, i].repeat(2, axis=0).repeat(2, axis=1)
normalized_nouns_maps[:, :, i] = (curr_noun_map - np.abs(curr_noun_map.min())) / curr_noun_map.max()
max_score = 0
all_scores = []
for c in range(self.num_segments):
cluster_mask = np.zeros_like(clusters)
cluster_mask[clusters == c] = 1
score_maps = [cluster_mask * normalized_nouns_maps[:, :, i] for i in range(len(nouns_indices))]
scores = [score_map.sum() / cluster_mask.sum() for score_map in score_maps]
all_scores.append(max(scores))
max_score = max(max(scores), max_score)
all_scores.remove(max_score)
mean_score = sum(all_scores) / len(all_scores)
for c in range(self.num_segments):
cluster_mask = np.zeros_like(clusters)
cluster_mask[clusters == c] = 1
score_maps = [cluster_mask * normalized_nouns_maps[:, :, i] for i in range(len(nouns_indices))]
scores = [score_map.sum() / cluster_mask.sum() for score_map in score_maps]
result[c] = self.nouns[np.argmax(np.array(scores))] if max(scores) > 1.4 * mean_score else "BG"
return result
def create_mask(self, clusters, cross_attention, attn_index):
cluster2noun = self.cluster2noun(clusters, cross_attention, attn_index)
mask = clusters.copy()
obj_segments = [c for c in cluster2noun if cluster2noun[c][1] in self.object_nouns]
for c in range(self.num_segments):
mask[clusters == c] = 1 if c in obj_segments else 0
return torch.from_numpy(mask).to("cuda")
def get_object_masks(self) -> Tuple[torch.Tensor]:
clusters_style_32, clusters_struct_32 = self.cluster(res=32)
clusters_style_64, clusters_struct_64 = self.cluster(res=64)
mask_style_32 = self.create_mask(clusters_style_32, self.cross_attention_32, STYLE_INDEX)
mask_struct_32 = self.create_mask(clusters_struct_32, self.cross_attention_32, STRUCT_INDEX)
mask_style_64 = self.create_mask(clusters_style_64, self.cross_attention_64, STYLE_INDEX)
mask_struct_64 = self.create_mask(clusters_struct_64, self.cross_attention_64, STRUCT_INDEX)
return mask_style_32, mask_struct_32, mask_style_64, mask_struct_64
|