Spaces:
Runtime error
Runtime error
| import nltk | |
| from sklearn.cluster import KMeans | |
| import numpy as np | |
| from src.attention_utils import aggregate_attention | |
| class Segmentor: | |
| def __init__(self, controller, prompts, num_segments, background_segment_threshold, res=32, background_nouns=[]): | |
| self.controller = controller | |
| self.prompts = prompts | |
| self.num_segments = num_segments | |
| self.background_segment_threshold = background_segment_threshold | |
| self.resolution = res | |
| self.background_nouns = background_nouns | |
| self.self_attention = aggregate_attention(controller, res=32, from_where=("up", "down"), prompts=prompts, | |
| is_cross=False, select=len(prompts) - 1) | |
| self.cross_attention = aggregate_attention(controller, res=16, from_where=("up", "down"), prompts=prompts, | |
| is_cross=True, select=len(prompts) - 1) | |
| tokenized_prompt = nltk.word_tokenize(prompts[-1]) | |
| self.nouns = [(i, word) for (i, (word, pos)) in enumerate(nltk.pos_tag(tokenized_prompt)) if pos[:2] == 'NN'] | |
| def __call__(self, *args, **kwargs): | |
| clusters = self.cluster() | |
| cluster2noun = self.cluster2noun(clusters) | |
| return cluster2noun | |
| def cluster(self): | |
| np.random.seed(1) | |
| resolution = self.self_attention.shape[0] | |
| attn = self.self_attention.cpu().numpy().reshape(resolution ** 2, resolution ** 2) | |
| kmeans = KMeans(n_clusters=self.num_segments, n_init=10).fit(attn) | |
| clusters = kmeans.labels_ | |
| clusters = clusters.reshape(resolution, resolution) | |
| return clusters | |
| def cluster2noun(self, clusters): | |
| result = {} | |
| nouns_indices = [index for (index, word) in self.nouns] | |
| nouns_maps = self.cross_attention.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() | |
| 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) > self.background_segment_threshold else "BG" | |
| return result | |
| def get_background_mask(self, obj_token_index): | |
| clusters = self.cluster() | |
| cluster2noun = self.cluster2noun(clusters) | |
| mask = clusters.copy() | |
| obj_segments = [c for c in cluster2noun if cluster2noun[c][0] == obj_token_index - 1] | |
| background_segments = [c for c in cluster2noun if cluster2noun[c] == "BG" or cluster2noun[c][1] in self.background_nouns] | |
| for c in range(self.num_segments): | |
| if c in background_segments and c not in obj_segments: | |
| mask[clusters == c] = 0 | |
| else: | |
| mask[clusters == c] = 1 | |
| return mask | |