from typing import Optional, Union, Tuple, List, Callable, Dict from tqdm.notebook import tqdm import torch import math from typing import List, Tuple, Union from PIL import Image import cv2 import numpy as np import os import re import torch from IPython.display import display from sklearn.cluster import KMeans import matplotlib.pyplot as plt from .ptp_utils import * import torchvision.transforms as transforms from sklearn.decomposition import PCA import pickle as pkl import torch.nn.functional as F import argparse from sklearn.metrics.cluster import adjusted_rand_score, normalized_mutual_info_score, fowlkes_mallows_score, v_measure_score transform_train = transforms.Compose([ transforms.ToPILImage(), ]) pca = PCA(n_components=3) def save_mask(mask, output_name): mask_image = transform_train(mask.float()) mask_image.save(output_name) def show_image(image): image = 255 * image / image.max() image = image.unsqueeze(-1).expand(*image.shape, 3) image = image.numpy().astype(np.uint8) image = np.array(Image.fromarray(image).resize((256, 256))) return image def cluster2noun_mod(clusters, background_segment_threshold, num_segments, nouns, cross_attention): REPEAT=clusters.shape[0]/cross_attention.shape[0] result = {} result_mask={} nouns_indices = [index for (index, word) in nouns] nouns_maps = cross_attention.cpu().numpy()[:, :, [i + 1 for i in nouns_indices]] nouns_maps = cross_attention.unsqueeze(-1).cpu().numpy() normalized_nouns_maps = np.zeros_like(nouns_maps).repeat(REPEAT, axis=0).repeat(REPEAT, axis=1) for i in range(nouns_maps.shape[-1]): curr_noun_map = nouns_maps[:, :, i].repeat(REPEAT, axis=0).repeat(REPEAT, axis=1) normalized_nouns_maps[:, :, i] = (curr_noun_map - np.abs(curr_noun_map.min())) / curr_noun_map.max() for c in range(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] = nouns[np.argmax(np.array(scores))] if max(scores) > background_segment_threshold else "BG" result_mask[c]=cluster_mask return result, result_mask def cluster2noun_(clusters, background_segment_threshold, num_segments, nouns, cross_attention): REPEAT=clusters.shape[0]/cross_attention.shape[0] result = {} result_mask={} nouns_indices = [index for (index, word) in nouns] nouns_maps = cross_attention.cpu().numpy()[:, :, [i + 1 for i in nouns_indices]] normalized_nouns_maps = np.zeros_like(nouns_maps).repeat(REPEAT, axis=0).repeat(REPEAT, axis=1) for i in range(nouns_maps.shape[-1]): curr_noun_map = nouns_maps[:, :, i].repeat(REPEAT, axis=0).repeat(REPEAT, axis=1) normalized_nouns_maps[:, :, i] = (curr_noun_map - np.abs(curr_noun_map.min())) / curr_noun_map.max() for c in range(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] = nouns[np.argmax(np.array(scores))] if max(scores) > background_segment_threshold else "BG" result_mask[c]=cluster_mask return result, result_mask def aggregate_attention( attention_maps, res: int, from_where: List[str], is_cross: bool, select: int, prompts,): out = [] num_pixels = res ** 2 for location in from_where: for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]: if item.shape[1] == num_pixels: cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select] out.append(cross_maps) out = torch.cat(out, dim=0) out = out.sum(0) / out.shape[0] return out.cpu(), attention_maps def cluster(self_attention, num_segments,): np.random.seed(1) frames = self_attention.shape[0] video_clusters = [] for i in range(frames): per_frame_attention = self_attention[i] print('per_frame_attention',per_frame_attention.shape) resolution, feat_dim = per_frame_attention.shape[0], per_frame_attention.shape[-1] attn = per_frame_attention.cpu().numpy().reshape(resolution ** 2, feat_dim) kmeans = KMeans(n_clusters=num_segments, n_init=10).fit(attn) clusters = kmeans.labels_ clusters = clusters.reshape(resolution, resolution) video_clusters.append(clusters) return video_clusters def run_clusters(avg_dict, resolution, dict_key, save_path, special_name, num_segments): video_clusters = cluster(avg_dict[dict_key][resolution], num_segments,) npy_name=f'cluster_{dict_key}_{resolution}_{special_name}.npy' np.save(os.path.join(save_path, npy_name), video_clusters) for i in range(len(video_clusters)): clusters = video_clusters[i] output_name=f'cluster_{dict_key}_{resolution}_{i}.png' plt.imshow(clusters) plt.axis('off') plt.savefig(os.path.join(save_path, output_name), bbox_inches='tight', pad_inches=0) def read_pkl(path,): with open(path,'rb') as f: dict_ = pkl.load(f) return dict_ def draw_pca(avg_dict, resolution, dict_key, save_path, special_name): RESOLUTION=resolution if avg_dict[dict_key][RESOLUTION].__len__() == 0: return before_pca = avg_dict[dict_key][RESOLUTION] frames = before_pca.shape[0] for i in range(frames): frame = before_pca[i] print('frame',frame.dtype) if isinstance(frame, torch.Tensor): frame = frame.reshape(RESOLUTION * RESOLUTION, -1).cpu().numpy() else: frame = frame.reshape(RESOLUTION * RESOLUTION, -1) pca.fit(frame) after_pca = pca.transform(frame) after_pca = after_pca.reshape(RESOLUTION,RESOLUTION,-1) pca_img_min = after_pca.min(axis=(0, 1)) pca_img_max = after_pca.max(axis=(0, 1)) pca_img = (after_pca - pca_img_min) / (pca_img_max - pca_img_min) output_name=f'pca_{dict_key}_{resolution}_{i}.png' pca_img = Image.fromarray((pca_img * 255).astype(np.uint8)) pca_img=pca_img.resize((512,512)) pca_img.save(os.path.join(save_path, output_name)) def image_normalize(numpy_array, save_path,output_name): numpy_array=numpy_array.cpu().numpy() img_min = numpy_array.min() img_max = numpy_array.max() normalize_array = (numpy_array - img_min) / (img_max - img_min) plt.imshow(normalize_array) plt.axis('off') plt.savefig(os.path.join(save_path, output_name), bbox_inches='tight', pad_inches=0) def cross_cosine_with_name(resolution, inv_avg_dict, denoise_avg_dict, indice, save_path, save_crossattn=False, noun_name = ''): inv_cross_attn = inv_avg_dict['attn'][resolution][:,:,indice] denoise_cross_attn = denoise_avg_dict['attn'][resolution][:,:,indice] if save_crossattn: image_normalize(inv_cross_attn, save_path, f'crossattn_{resolution}_inv_{noun_name}.png') image_normalize(denoise_cross_attn, save_path, f'crossattn_{resolution}_denoise_{noun_name}.png') return F.cosine_similarity(inv_cross_attn.reshape(1,-1), denoise_cross_attn.reshape(1,-1)) def cross_cosine(resolution, inv_avg_dict, denoise_avg_dict, indice, save_path, save_crossattn=False,): inv_cross_attn = inv_avg_dict['attn'][resolution][:,:,indice] denoise_cross_attn = denoise_avg_dict['attn'][resolution][:,:,indice] if save_crossattn: image_normalize(inv_cross_attn, save_path, f'crossattn_{resolution}_inv.png') image_normalize(denoise_cross_attn, save_path, f'crossattn_{resolution}_denoise.png') return F.cosine_similarity(inv_cross_attn.reshape(1,-1), denoise_cross_attn.reshape(1,-1)) def save_crossattn(input_path, caption, inv_cross_avg_dict, denoise_cross_avg_dict, results_folder, RES=16): org_image = Image.open(input_path).convert("RGB") prompts=["<|startoftext|>",] + caption.split(' ') + ["<|endoftext|>",] inv_crossattn = inv_cross_avg_dict['attn'][RES] denoise_crossattn = denoise_cross_avg_dict['attn'][RES] attn_img1, mask_img1, _ = show_cross_attention_plus_orig_img(prompts, inv_crossattn, orig_image=org_image) attn_img2, mask_img2, _ = show_cross_attention_plus_orig_img(prompts, denoise_crossattn, orig_image=org_image) attn_img1.save(os.path.join(results_folder,'crossattn_inv.png')) attn_img2.save(os.path.join(results_folder,'crossattn_denoise.png')) mask_img1.save(os.path.join(results_folder,'crossattn_inv_mask.png')) mask_img2.save(os.path.join(results_folder,'crossattn_denoise_mask.png'))