Spaces:
Runtime error
Runtime error
File size: 3,293 Bytes
c4e6a63 |
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 |
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
|