import os
import clip
import torch
import open_clip
import numpy as np
from torchvision.datasets import CIFAR100
from tqdm import tqdm
import torchvision.transforms as transforms
import warnings
with warnings.catch_warnings():
warnings.simplefilter(action='ignore', category=FutureWarning)
with warnings.catch_warnings():
warnings.simplefilter(action='ignore', category=UserWarning)
import torchvision
import pandas as pd
from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import pickle
class FACET(Dataset):
"""Face Landmarks dataset."""
def __init__(self, paths, labels, root_dir, file_extension=".jpg", transform=None):
"""
Arguments:
csv_file (string): Path to the csv file with annotations.
root_dir (string): Directory with all the images.
transform (callable, optional): Optional transform to be applied
on a sample.
"""
self.fpaths = paths
self.extension = file_extension
self.labels = labels
self.root_dir = root_dir
self.transform = transform
def __len__(self):
return len(self.fpaths)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
img_name = os.path.join(self.root_dir,
str(self.fpaths[idx])+self.extension)
image = self.transform(Image.open(img_name).convert('RGB'))
label = self.labels[idx]
return image, label
imagenet_templates = [
'a bad photo of a {}.',
'a photo of many {}.',
'a sculpture of a {}.',
'a photo of the hard to see {}.',
'a low resolution photo of the {}.',
'a rendering of a {}.',
'graffiti of a {}.',
'a bad photo of the {}.',
'a cropped photo of the {}.',
'a tattoo of a {}.',
'the embroidered {}.',
'a photo of a hard to see {}.',
'a bright photo of a {}.',
'a photo of a clean {}.',
'a photo of a dirty {}.',
'a dark photo of the {}.',
'a drawing of a {}.',
'a photo of my {}.',
'the plastic {}.',
'a photo of the cool {}.',
'a close-up photo of a {}.',
'a black and white photo of the {}.',
'a painting of the {}.',
'a painting of a {}.',
'a pixelated photo of the {}.',
'a sculpture of the {}.',
'a bright photo of the {}.',
'a cropped photo of a {}.',
'a plastic {}.',
'a photo of the dirty {}.',
'a jpeg corrupted photo of a {}.',
'a blurry photo of the {}.',
'a photo of the {}.',
'a good photo of the {}.',
'a rendering of the {}.',
'a {} in a video game.',
'a photo of one {}.',
'a doodle of a {}.',
'a close-up photo of the {}.',
'a photo of a {}.',
'the origami {}.',
'the {} in a video game.',
'a sketch of a {}.',
'a doodle of the {}.',
'a origami {}.',
'a low resolution photo of a {}.',
'the toy {}.',
'a rendition of the {}.',
'a photo of the clean {}.',
'a photo of a large {}.',
'a rendition of a {}.',
'a photo of a nice {}.',
'a photo of a weird {}.',
'a blurry photo of a {}.',
'a cartoon {}.',
'art of a {}.',
'a sketch of the {}.',
'a embroidered {}.',
'a pixelated photo of a {}.',
'itap of the {}.',
'a jpeg corrupted photo of the {}.',
'a good photo of a {}.',
'a plushie {}.',
'a photo of the nice {}.',
'a photo of the small {}.',
'a photo of the weird {}.',
'the cartoon {}.',
'art of the {}.',
'a drawing of the {}.',
'a photo of the large {}.',
'a black and white photo of a {}.',
'the plushie {}.',
'a dark photo of a {}.',
'itap of a {}.',
'graffiti of the {}.',
'a toy {}.',
'itap of my {}.',
'a photo of a cool {}.',
'a photo of a small {}.',
'a tattoo of the {}.',
]
models = (
# CLIP OpenAI
"ViT-B/16",
"ViT-B/32",
"ViT-L/14",
"RN50",
"RN101",
# CLIP OpenCLIP
"vit_b_16_400m",
"vit_b_16_2b",
"vit_l_14_400m",
"vit_l_14_2b",
"vit_b_32_400m",
"vit_b_32_2b",
)
weights = (
# CLIP OpenAI
"OpenAI hub",
"OpenAI hub",
"OpenAI hub",
"OpenAI hub",
"OpenAI hub",
# CLIP OpenCLIP
"OpenCLIP hub",
"OpenCLIP hub",
"OpenCLIP hub",
"OpenCLIP hub",
"OpenCLIP hub",
"OpenCLIP hub",
)
facet_annotations_file_path = "INSERT_HERE/annotations.csv"
facet_root = "?????" # where the in-painted images are stored, the following structure is expected:
# facet_root/
# facet_paper_skin_ours/
# facet_paper_clothes_only/
# facet_paper_skin_ours_occupation_prompt/
# facet_paper_clothes_only_occupation_prompt/
# facet_paper_whole_body/
# facet_paper_whole_body_occupation_prompt/
facet = pd.read_csv(facet_annotations_file_path, header=0).rename(columns={'Unnamed: 0': 'sample_idx'})
device = "cuda" if torch.cuda.is_available() else "cpu"
experiments = ["facet_paper_skin_ours", "facet_paper_clothes_only", "facet_paper_skin_ours_occupation_prompt", "facet_paper_clothes_only_occupation_prompt",
"facet_paper_whole_body", "facet_paper_whole_body_occupation_prompt"
]
for experiment in experiments:
for model_name, weight in zip(models, weights):
print( "\n\n",model_name, experiment)
preprocess = None
if model_name == "vit_b_16_400m":
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-16', pretrained='laion400m_e32')
clip_src = "OpenCLIP"
elif model_name == "vit_b_16_2b":
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-16', pretrained='laion2b_s34b_b88k')
clip_src = "OpenCLIP"
elif model_name == "vit_b_32_400m":
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion400m_e32')
clip_src = "OpenCLIP"
elif model_name == "vit_b_32_2b":
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
clip_src = "OpenCLIP"
elif model_name == "vit_l_14_400m":
model, _, preprocess = open_clip.create_model_and_transforms('ViT-L-14', pretrained='laion400m_e32')
clip_src = "OpenCLIP"
elif model_name == "vit_l_14_2b":
model, _, preprocess = open_clip.create_model_and_transforms('ViT-L-14', pretrained='laion2b_s32b_b82k')
clip_src = "OpenCLIP"
elif "ViT" in model_name:
model, preprocess = clip.load(model_name, device)
clip_src = "OpenAI"
elif "RN" in model_name:
model, preprocess = clip.load(model_name, device)
clip_src = "OpenAI"
else:
raise NotImplementedError
model.cuda()
occupations = ['backpacker', 'ballplayer', 'bartender', 'basketball_player', 'boatman', 'carpenter', 'cheerleader', 'climber', 'computer_user', 'craftsman', 'dancer', 'disk_jockey', 'doctor', 'drummer', 'electrician', 'farmer', 'fireman', 'flutist', 'gardener', 'guard', 'guitarist', 'gymnast', 'hairdresser', 'horseman', 'judge', 'laborer', 'lawman', 'lifeguard', 'machinist', 'motorcyclist', 'nurse', 'painter', 'patient', 'prayer', 'referee', 'repairman', 'reporter', 'retailer', 'runner', 'sculptor', 'seller', 'singer', 'skateboarder', 'soccer_player', 'soldier', 'speaker', 'student', 'teacher', 'tennis_player', 'trumpeter', 'waiter']
tokens_occupations = clip.tokenize([f"A photo of a {occupation}" for occupation in occupations]).cuda()
facet_img_root = facet_save_root + "/" experiment + "/"
out_dir = experiment + "_zero_shot"
if not os.path.exists(out_dir):
os.makedirs(out_dir)
fnames = list(os.listdir(facet_img_root))
for attribute_value in ["only_original_male", "only_original_female", "original", "male_to_female", "male_to_male", "female_to_female", "female_to_male"]:
print(f"----{attribute_value}----")
facet = pd.read_csv("../../datasets/facet/annotations/annotations.csv", header=0).rename(columns={'Unnamed: 0': 'sample_idx'}) # Bounding boxes
extension = ".png"
processed_synthetic_samples = set()
for fname in fnames:
bbid, target_attr = fname.split("_")[0], "_".join(fname.split("_")[1:]).split(".")[0]
if "only" in attribute_value:
if target_attr=="original" and bbid not in processed_synthetic_samples:
processed_synthetic_samples.add(int(bbid))
elif target_attr==attribute_value and bbid not in processed_synthetic_samples:
processed_synthetic_samples.add(int(bbid))
if attribute_value == "only_original_male":
facet = facet[facet.person_id.isin(processed_synthetic_samples)]
facet = facet[facet.gender_presentation_na != 1]
facet = facet[facet.gender_presentation_non_binary != 1]
facet = facet[(facet.gender_presentation_masc == 1)]
elif attribute_value == "only_original_female":
facet = facet[facet.person_id.isin(processed_synthetic_samples)]
facet = facet[facet.gender_presentation_na != 1]
facet = facet[facet.gender_presentation_non_binary != 1]
facet = facet[(facet.gender_presentation_fem == 1)]
else:
facet = facet[facet.person_id.isin(processed_synthetic_samples)]
facet = facet[facet.gender_presentation_na != 1]
facet = facet[facet.gender_presentation_non_binary != 1]
facet = facet[(facet.gender_presentation_masc == 1) | (facet.gender_presentation_fem == 1)]
facet["class1"] = facet["class1"].apply(lambda val: int(occupations.index(val)))
bsize = 512
predictions = []
acc = my_acc = 0
n_batches = 0
def zeroshot_classifier(classnames, templates):
with torch.no_grad():
zeroshot_weights = []
for classname in tqdm(classnames):
texts = [template.format(classname) for template in templates] #format with class
texts = clip.tokenize(texts).cuda() #tokenize
class_embeddings = model.encode_text(texts) #embed with text encoder
class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
class_embedding = class_embeddings.mean(dim=0)
class_embedding /= class_embedding.norm()
zeroshot_weights.append(class_embedding)
zeroshot_weights = torch.stack(zeroshot_weights, dim=1).cuda()
return zeroshot_weights
if "only" in attribute_value:
dataset = FACET(facet.person_id.values, torch.tensor(facet.class1.values), facet_img_root, transform=preprocess, file_extension=f"_original.png")
else:
dataset = FACET(facet.person_id.values, torch.tensor(facet.class1.values), facet_img_root, transform=preprocess, file_extension=f"_{attribute_value}.png")
dataloader = DataLoader(dataset, batch_size=bsize, shuffle=False, num_workers=6, drop_last=False,)
zeroshot_weights = zeroshot_classifier(occupations[:39], imagenet_templates)
for imgs, labels in tqdm(dataloader):
with torch.no_grad(), torch.cuda.amp.autocast():
if clip_src == "OpenAI":
# CLIP
image_features = model.encode_image(imgs.half().cuda())
image_features /= image_features.norm(dim=-1, keepdim=True)
logits = 100. * image_features @ zeroshot_weights
probs = logits.softmax(dim=-1).cpu().numpy()
else:
# OpenCLIP
image_features = model.encode_image(imgs.half().cuda())
image_features /= image_features.norm(dim=-1, keepdim=True)
probs = (100. * image_features @ zeroshot_weights).softmax(dim=-1).cpu().numpy()
preds_batch = np.argmax(probs, axis=-1)
predictions += preds_batch.tolist()
acc += torch.sum(torch.tensor(preds_batch).cuda()==labels.cuda()) / preds_batch.shape[0]
n_batches += 1
print(model_name, "acc: ", acc / n_batches, "%")
results = pd.DataFrame({"person_id": facet.person_id.values,
"inpainted_attribute": attribute_value,
"age_presentation_young": facet.age_presentation_young.values,
"age_presentation_middle": facet.age_presentation_middle.values,
"age_presentation_older": facet.age_presentation_older.values,
"gender_presentation_fem": facet.gender_presentation_fem.values,
"gender_presentation_masc": facet.gender_presentation_masc.values,
"gt_class_label": facet.class1.values,
"class_predictions": predictions
})
results.to_csv(f'{out_dir}/{model_name.replace("/", "_").replace("-", "_")}_{attribute_value}_predictions.csv')
with open(f'{out_dir}/{model_name.replace("/", "_").replace("-", "_")}_{attribute_value}_accuracy.txt', "w") as of:
of.write(str((acc/n_batches).item()))
with open(f'{out_dir}/{model_name.replace("/", "_").replace("-", "_")}_{attribute_value}.pkl', "wb") as f:
pickle.dump(predictions, f)