video_diffusion/prompt_attention/sd_study_utils.py

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'))