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, attention_threshold=0.2): REPEAT = clusters.shape[0] // cross_attention.shape[0] result = {} result_mask = {} print('cross_attention',cross_attention.shape) # 提取名词索引和对应的注意力图 nouns_indices = [index for (index, word) in nouns] nouns_maps = cross_attention.cpu().numpy()[:, :, [i + 1 for i in nouns_indices]] print('nouns_maps', nouns_maps.shape) normalized_nouns_maps = nouns_maps #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_map = (curr_noun_map - np.abs(curr_noun_map.min())) / curr_noun_map.max() # # 应用阈值，将低于阈值的部分设为 0 # #normalized_map[normalized_map < attention_threshold] = 0 # normalized_nouns_maps[:, :, i] = normalized_map print('normalized_nouns_maps', normalized_nouns_maps.shape) #show_normalized_nouns_maps(normalized_nouns_maps, nouns, logdir) # 用于记录已经分配的单词 assigned_nouns = set() 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] # 找出最高分的名词，并确保未被分配过 sorted_scores_indices = np.argsort(scores)[::-1] assigned_word = None for idx in sorted_scores_indices: if scores[idx] > background_segment_threshold and nouns[idx] not in assigned_nouns: assigned_word = nouns[idx] assigned_nouns.add(nouns[idx]) # 记录这个单词已分配 break # 如果没有找到合适的名词，强制分配最高分的未分配名词 if assigned_word is None and len(sorted_scores_indices) > 0: for idx in sorted_scores_indices: if nouns[idx] not in assigned_nouns: assigned_word = nouns[idx] assigned_nouns.add(nouns[idx]) # 记录这个单词已分配 break if assigned_word: result[c] = assigned_word 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'))