#!/usr/bin/env python # coding=utf-8 # Copyright 2023 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and import argparse import hashlib import logging import math import os import warnings from pathlib import Path from functools import reduce import numpy as np import torch import torch.nn.functional as F import torch.utils.checkpoint import transformers from packaging import version from PIL import Image from torch.utils.data import Dataset, DataLoader from torchvision import transforms from tqdm.auto import tqdm from transformers import AutoTokenizer, PretrainedConfig, ViTFeatureExtractor, ViTModel import lpips import json from PIL import Image import requests from transformers import AutoProcessor, AutoTokenizer, CLIPModel import torchvision.transforms.functional as TF from torch.nn.functional import cosine_similarity from torchvision.transforms import Compose, ToTensor, Normalize, Resize, ToPILImage import re def get_prompt(subject_name, prompt_idx): subject_names = [ "backpack", "backpack_dog", "bear_plushie", "berry_bowl", "can", "candle", "cat", "cat2", "clock", "colorful_sneaker", "dog", "dog2", "dog3", "dog5", "dog6", "dog7", "dog8", "duck_toy", "fancy_boot", "grey_sloth_plushie", "monster_toy", "pink_sunglasses", "poop_emoji", "rc_car", "red_cartoon", "robot_toy", "shiny_sneaker", "teapot", "vase", "wolf_plushie", ] class_tokens = [ "backpack", "backpack", "stuffed animal", "bowl", "can", "candle", "cat", "cat", "clock", "sneaker", "dog", "dog", "dog", "dog", "dog", "dog", "dog", "toy", "boot", "stuffed animal", "toy", "glasses", "toy", "toy", "cartoon", "toy", "sneaker", "teapot", "vase", "stuffed animal", ] class_token = class_tokens[subject_names.index(subject_name)] prompt_list = [ f"a qwe {class_token} in the jungle", f"a qwe {class_token} in the snow", f"a qwe {class_token} on the beach", f"a qwe {class_token} on a cobblestone street", f"a qwe {class_token} on top of pink fabric", f"a qwe {class_token} on top of a wooden floor", f"a qwe {class_token} with a city in the background", f"a qwe {class_token} with a mountain in the background", f"a qwe {class_token} with a blue house in the background", f"a qwe {class_token} on top of a purple rug in a forest", f"a qwe {class_token} wearing a red hat", f"a qwe {class_token} wearing a santa hat", f"a qwe {class_token} wearing a rainbow scarf", f"a qwe {class_token} wearing a black top hat and a monocle", f"a qwe {class_token} in a chef outfit", f"a qwe {class_token} in a firefighter outfit", f"a qwe {class_token} in a police outfit", f"a qwe {class_token} wearing pink glasses", f"a qwe {class_token} wearing a yellow shirt", f"a qwe {class_token} in a purple wizard outfit", f"a red qwe {class_token}", f"a purple qwe {class_token}", f"a shiny qwe {class_token}", f"a wet qwe {class_token}", f"a cube shaped qwe {class_token}", ] return prompt_list[int(prompt_idx)] class PromptDatasetCLIP(Dataset): def __init__(self, subject_name, data_dir_B, tokenizer, processor, epoch=None): self.data_dir_B = data_dir_B subject_name, prompt_idx = subject_name.split('-') data_dir_B = os.path.join(self.data_dir_B, str(epoch)) self.image_lst = [os.path.join(data_dir_B, f) for f in os.listdir(data_dir_B) if f.endswith(".png")] self.prompt_lst = [get_prompt(subject_name, prompt_idx)] * len(self.image_lst) self.tokenizer = tokenizer self.processor = processor self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def __len__(self): return len(self.image_lst) def __getitem__(self, idx): image_path = self.image_lst[idx] image = Image.open(image_path) prompt = self.prompt_lst[idx] extrema = image.getextrema() if all(min_val == max_val == 0 for min_val, max_val in extrema): return None, None else: prompt_inputs = self.tokenizer([prompt], padding=True, return_tensors="pt") image_inputs = self.processor(images=image, return_tensors="pt") return image_inputs, prompt_inputs class PairwiseImageDatasetCLIP(Dataset): def __init__(self, subject_name, data_dir_A, data_dir_B, processor, epoch): self.data_dir_A = data_dir_A self.data_dir_B = data_dir_B subject_name, prompt_idx = subject_name.split('-') self.data_dir_A = os.path.join(self.data_dir_A, subject_name) self.image_files_A = [os.path.join(self.data_dir_A, f) for f in os.listdir(self.data_dir_A) if f.endswith(".jpg")] data_dir_B = os.path.join(self.data_dir_B, str(epoch)) self.image_files_B = [os.path.join(data_dir_B, f) for f in os.listdir(data_dir_B) if f.endswith(".png")] self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.processor = processor def __len__(self): return len(self.image_files_A) * len(self.image_files_B) def __getitem__(self, index): index_A = index // len(self.image_files_B) index_B = index % len(self.image_files_B) image_A = Image.open(self.image_files_A[index_A]) # .convert("RGB") image_B = Image.open(self.image_files_B[index_B]) # .convert("RGB") extrema_A = image_A.getextrema() extrema_B = image_B.getextrema() if all(min_val == max_val == 0 for min_val, max_val in extrema_A) or all(min_val == max_val == 0 for min_val, max_val in extrema_B): return None, None else: inputs_A = self.processor(images=image_A, return_tensors="pt") inputs_B = self.processor(images=image_B, return_tensors="pt") return inputs_A, inputs_B class PairwiseImageDatasetDINO(Dataset): def __init__(self, subject_name, data_dir_A, data_dir_B, feature_extractor, epoch): self.data_dir_A = data_dir_A self.data_dir_B = data_dir_B subject_name, prompt_idx = subject_name.split('-') self.data_dir_A = os.path.join(self.data_dir_A, subject_name) self.image_files_A = [os.path.join(self.data_dir_A, f) for f in os.listdir(self.data_dir_A) if f.endswith(".jpg")] data_dir_B = os.path.join(self.data_dir_B, str(epoch)) self.image_files_B = [os.path.join(data_dir_B, f) for f in os.listdir(data_dir_B) if f.endswith(".png")] self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.feature_extractor = feature_extractor def __len__(self): return len(self.image_files_A) * len(self.image_files_B) def __getitem__(self, index): index_A = index // len(self.image_files_B) index_B = index % len(self.image_files_B) image_A = Image.open(self.image_files_A[index_A]) # .convert("RGB") image_B = Image.open(self.image_files_B[index_B]) # .convert("RGB") extrema_A = image_A.getextrema() extrema_B = image_B.getextrema() if all(min_val == max_val == 0 for min_val, max_val in extrema_A) or all(min_val == max_val == 0 for min_val, max_val in extrema_B): return None, None else: inputs_A = self.feature_extractor(images=image_A, return_tensors="pt") inputs_B = self.feature_extractor(images=image_B, return_tensors="pt") return inputs_A, inputs_B class PairwiseImageDatasetLPIPS(Dataset): def __init__(self, subject_name, data_dir_A, data_dir_B, epoch): self.data_dir_A = data_dir_A self.data_dir_B = data_dir_B subject_name, prompt_idx = subject_name.split('-') self.data_dir_A = os.path.join(self.data_dir_A, subject_name) self.image_files_A = [os.path.join(self.data_dir_A, f) for f in os.listdir(self.data_dir_A) if f.endswith(".jpg")] data_dir_B = os.path.join(self.data_dir_B, str(epoch)) self.image_files_B = [os.path.join(data_dir_B, f) for f in os.listdir(data_dir_B) if f.endswith(".png")] self.transform = Compose([ Resize((512, 512)), ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def __len__(self): return len(self.image_files_A) * len(self.image_files_B) def __getitem__(self, index): index_A = index // len(self.image_files_B) index_B = index % len(self.image_files_B) image_A = Image.open(self.image_files_A[index_A]) # .convert("RGB") image_B = Image.open(self.image_files_B[index_B]) # .convert("RGB") extrema_A = image_A.getextrema() extrema_B = image_B.getextrema() if all(min_val == max_val == 0 for min_val, max_val in extrema_A) or all(min_val == max_val == 0 for min_val, max_val in extrema_B): return None, None else: if self.transform: image_A = self.transform(image_A) image_B = self.transform(image_B) return image_A, image_B def clip_text(subject_name, image_dir): criterion = 'clip_text' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(device) # Get the text features tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-large-patch14") # Get the image features processor = AutoProcessor.from_pretrained("openai/clip-vit-large-patch14") epochs = sorted([int(epoch) for epoch in os.listdir(image_dir)]) best_mean_similarity = 0 mean_similarity_list = [] for epoch in epochs: similarity = [] dataset = PromptDatasetCLIP(subject_name, image_dir, tokenizer, processor, epoch) dataloader = DataLoader(dataset, batch_size=32) for i in range(len(dataset)): image_inputs, prompt_inputs = dataset[i] if image_inputs is not None and prompt_inputs is not None: image_inputs['pixel_values'] = image_inputs['pixel_values'].to(device) prompt_inputs['input_ids'] = prompt_inputs['input_ids'].to(device) prompt_inputs['attention_mask'] = prompt_inputs['attention_mask'].to(device) # print(prompt_inputs) image_features = model.get_image_features(**image_inputs) text_features = model.get_text_features(**prompt_inputs) sim = cosine_similarity(image_features, text_features) #image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True) #text_features = text_features / text_features.norm(p=2, dim=-1, keepdim=True) #logit_scale = model.logit_scale.exp() #sim = torch.matmul(text_features, image_features.t()) * logit_scale similarity.append(sim.item()) if similarity: mean_similarity = torch.tensor(similarity).mean().item() mean_similarity_list.append(mean_similarity) best_mean_similarity = max(best_mean_similarity, mean_similarity) print(f'epoch: {epoch}, criterion: {criterion}, mean_similarity: {mean_similarity}({best_mean_similarity})') else: mean_similarity_list.append(0) print(f'epoch: {epoch}, criterion: {criterion}, mean_similarity: {0}({best_mean_similarity})') return mean_similarity_list def clip_image(subject_name, image_dir, dreambooth_dir='dreambooth/dataset'): criterion = 'clip_image' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(device) # Get the image features processor = AutoProcessor.from_pretrained("openai/clip-vit-large-patch14") epochs = sorted([int(epoch) for epoch in os.listdir(image_dir)]) best_mean_similarity = 0 mean_similarity_list = [] for epoch in epochs: similarity = [] dataset = PairwiseImageDatasetCLIP(subject_name, dreambooth_dir, image_dir, processor, epoch) # dataset = SelfPairwiseImageDatasetCLIP(subject, './data', processor) for i in range(len(dataset)): inputs_A, inputs_B = dataset[i] if inputs_A is not None and inputs_B is not None: inputs_A['pixel_values'] = inputs_A['pixel_values'].to(device) inputs_B['pixel_values'] = inputs_B['pixel_values'].to(device) image_A_features = model.get_image_features(**inputs_A) image_B_features = model.get_image_features(**inputs_B) image_A_features = image_A_features / image_A_features.norm(p=2, dim=-1, keepdim=True) image_B_features = image_B_features / image_B_features.norm(p=2, dim=-1, keepdim=True) logit_scale = model.logit_scale.exp() sim = torch.matmul(image_A_features, image_B_features.t()) # * logit_scale similarity.append(sim.item()) if similarity: mean_similarity = torch.tensor(similarity).mean().item() best_mean_similarity = max(best_mean_similarity, mean_similarity) mean_similarity_list.append(mean_similarity) print(f'epoch: {epoch}, criterion: {criterion}, mean_similarity: {mean_similarity}({best_mean_similarity})') else: mean_similarity_list.append(0) print(f'epoch: {epoch}, criterion: {criterion}, mean_similarity: {0}({best_mean_similarity})') return mean_similarity_list def dino(subject_name, image_dir, dreambooth_dir='dreambooth/dataset'): criterion = 'dino' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = ViTModel.from_pretrained('facebook/dino-vits16').to(device) feature_extractor = ViTFeatureExtractor.from_pretrained('facebook/dino-vits16') epochs = sorted([int(epoch) for epoch in os.listdir(image_dir)]) best_mean_similarity = 0 mean_similarity_list = [] for epoch in epochs: similarity = [] # dataset = PairwiseImageDatasetDINO(subject, './data', image_dir, feature_extractor, epoch) dataset = PairwiseImageDatasetDINO(subject_name, dreambooth_dir, image_dir, feature_extractor, epoch) # dataset = SelfPairwiseImageDatasetDINO(subject, './data', feature_extractor) for i in range(len(dataset)): inputs_A, inputs_B = dataset[i] if inputs_A is not None and inputs_B is not None: inputs_A['pixel_values'] = inputs_A['pixel_values'].to(device) inputs_B['pixel_values'] = inputs_B['pixel_values'].to(device) outputs_A = model(**inputs_A) image_A_features = outputs_A.last_hidden_state[:, 0, :] outputs_B = model(**inputs_B) image_B_features = outputs_B.last_hidden_state[:, 0, :] image_A_features = image_A_features / image_A_features.norm(p=2, dim=-1, keepdim=True) image_B_features = image_B_features / image_B_features.norm(p=2, dim=-1, keepdim=True) sim = torch.matmul(image_A_features, image_B_features.t()) # * logit_scale similarity.append(sim.item()) mean_similarity = torch.tensor(similarity).mean().item() best_mean_similarity = max(best_mean_similarity, mean_similarity) mean_similarity_list.append(mean_similarity) print(f'epoch: {epoch}, criterion: {criterion}, mean_similarity: {mean_similarity}({best_mean_similarity})') return mean_similarity_list def lpips_image(subject_name, image_dir, dreambooth_dir='dreambooth/dataset'): criterion = 'lpips_image' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Set up the LPIPS model (vgg=True uses the VGG-based model from the paper) loss_fn = lpips.LPIPS(net='vgg').to(device) # 有可能有些epoch没跑全 epochs = sorted([int(epoch) for epoch in os.listdir(image_dir)]) mean_similarity_list = [] best_mean_similarity = 0 for epoch in epochs: similarity = [] dataset = PairwiseImageDatasetLPIPS(subject_name, dreambooth_dir, image_dir, epoch) # dataset = SelfPairwiseImageDatasetLPIPS(subject, './data') for i in range(len(dataset)): image_A, image_B = dataset[i] if image_A is not None and image_B is not None: image_A = image_A.to(device) image_B = image_B.to(device) # Calculate LPIPS between the two images distance = loss_fn(image_A, image_B) similarity.append(distance.item()) mean_similarity = torch.tensor(similarity).mean().item() best_mean_similarity = max(best_mean_similarity, mean_similarity) mean_similarity_list.append(mean_similarity) print(f'epoch: {epoch}, criterion: LPIPS distance, mean_similarity: {mean_similarity}({best_mean_similarity})') return mean_similarity_list if __name__ == "__main__": image_dir = 'log_hra/lr_1e-4_r_8/' subject_dirs, subject_names = [], [] for name in os.listdir(image_dir): if os.path.isdir(os.path.join(image_dir, name)): subject_dirs.append(os.path.join(image_dir, name)) subject_names.append(name) results_path = os.path.join(image_dir, 'true_results.json') # {'backpack-0':{'DINO':[x, ...], 'CLIP-I':[x, ...], 'CLIP-T':[x, ...], 'LPIPS':[x, ...],}} results_dict = dict() if os.path.exists(results_path): with open(results_path, 'r') as f: results = f.__iter__() while True: try: result_json = json.loads(next(results)) results_dict.update(result_json) except StopIteration: print("finish extraction.") break for idx in range(len(subject_names)): subject_name = subject_names[idx] subject_dir = subject_dirs[idx] if subject_name in results_dict: continue print(f'evaluating {subject_dir}') dino_sim = dino(subject_name, subject_dir) clip_i_sim = clip_image(subject_name, subject_dir) clip_t_sim = clip_text(subject_name, subject_dir) lpips_sim = lpips_image(subject_name, subject_dir) subject_result = {'DINO': dino_sim, 'CLIP-I': clip_i_sim, 'CLIP-T': clip_t_sim, 'LPIPS': lpips_sim} print(subject_result) with open(results_path,'a') as f: json_string = json.dumps({subject_name: subject_result}) f.write(json_string + "\n")