# Authors: Hui Ren (rhfeiyang.github.io) from transformers import CLIPProcessor, CLIPModel import torch import numpy as np import os os.environ["TOKENIZERS_PARALLELISM"] = "false" from tqdm import tqdm class Caption_filter: def __init__(self, filter_prompts=["painting", "paintings", "art", "artwork", "drawings", "sketch", "sketches", "illustration", "illustrations", "sculpture","sculptures", "installation", "printmaking", "digital art", "conceptual art", "mosaic", "tapestry", "abstract", "realism", "surrealism", "impressionism", "expressionism", "cubism", "minimalism", "baroque", "rococo", "pop art", "art nouveau", "art deco", "futurism", "dadaism", "stamp", "stamps", "advertisement", "advertisements","logo", "logos" ],): self.filter_prompts = filter_prompts self.total_count=0 self.filter_count=[0]*len(filter_prompts) def reset(self): self.total_count=0 self.filter_count=[0]*len(self.filter_prompts) def filter(self, captions): filter_result = [] for caption in captions: words = caption[0] if words == None: filter_result.append((True, "None")) continue words = words.lower() words = words.split() filt = False reason=None for i, filter_keyword in enumerate(self.filter_prompts): key_len = len(filter_keyword.split()) for j in range(len(words)-key_len+1): if " ".join(words[j:j+key_len]) == filter_keyword: self.filter_count[i] += 1 filt = True reason = filter_keyword break if filt: break filter_result.append((filt, reason)) self.total_count += 1 return filter_result class Clip_filter: prompt_threshold = { "painting": 17, "art": 17.5, "artwork": 19, "drawing": 15.8, "sketch": 17, "illustration": 15, "sculpture": 19.2, "installation art": 20, "printmaking art": 16.3, "digital art": 15, "conceptual art": 18, "mosaic art": 19, "tapestry": 16, "abstract art":16.5, "realism art": 16, "surrealism art": 15, "impressionism art": 17, "expressionism art": 17, "cubism art": 15, "minimalism art": 16, "baroque art": 17.5, "rococo art": 17, "pop art": 16, "art nouveau": 19, "art deco": 19, "futurism art": 16.5, "dadaism art": 16.5, "stamp": 18, "advertisement": 16.5, "logo": 15.5, } @torch.no_grad() def __init__(self, positive_prompt=["painting", "art", "artwork", "drawing", "sketch", "illustration", "sculpture", "installation art", "printmaking art", "digital art", "conceptual art", "mosaic art", "tapestry", "abstract art", "realism art", "surrealism art", "impressionism art", "expressionism art", "cubism art", "minimalism art", "baroque art", "rococo art", "pop art", "art nouveau", "art deco", "futurism art", "dadaism art", "stamp", "advertisement", "logo" ], device="cuda"): self.device = device self.model = (CLIPModel.from_pretrained("openai/clip-vit-large-patch14")).to(device) self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14") self.positive_prompt = positive_prompt self.text = self.positive_prompt self.tokenizer = self.processor.tokenizer self.image_processor = self.processor.image_processor self.text_encoding = self.tokenizer(self.text, return_tensors="pt", padding=True).to(device) self.text_features = self.model.get_text_features(**self.text_encoding) self.text_features = self.text_features / self.text_features.norm(p=2, dim=-1, keepdim=True) @torch.no_grad() def similarity(self, image): # inputs = self.processor(text=self.text, images=image, return_tensors="pt", padding=True) image_processed = self.image_processor(image, return_tensors="pt", padding=True).to(self.device, non_blocking=True) inputs = {**self.text_encoding, **image_processed} outputs = self.model(**inputs) logits_per_image = outputs.logits_per_image return logits_per_image def get_logits(self, image): logits_per_image = self.similarity(image) return logits_per_image.cpu() def get_image_features(self, image): image_processed = self.image_processor(image, return_tensors="pt", padding=True).to(self.device, non_blocking=True) image_features = self.model.get_image_features(**image_processed) return image_features class Art_filter: def __init__(self): self.caption_filter = Caption_filter() self.clip_filter = Clip_filter() def caption_filt(self, dataloader): self.caption_filter.reset() dataloader.dataset.get_img = False dataloader.dataset.get_cap = True remain_ids = [] filtered_ids = [] for i, batch in tqdm(enumerate(dataloader)): captions = batch["text"] filter_result = self.caption_filter.filter(captions) for j, (filt, reason) in enumerate(filter_result): if filt: filtered_ids.append((batch["ids"][j], reason)) if i%10==0: print(f"Filtered caption: {captions[j]}, reason: {reason}") else: remain_ids.append(batch["ids"][j]) return {"remain_ids":remain_ids, "filtered_ids":filtered_ids, "total_count":self.caption_filter.total_count, "filter_count":self.caption_filter.filter_count, "filter_prompts":self.caption_filter.filter_prompts} def clip_filt(self, clip_logits_ckpt:dict): logits = clip_logits_ckpt["clip_logits"] ids = clip_logits_ckpt["ids"] text = clip_logits_ckpt["text"] filt_mask = torch.zeros(logits.shape[0], dtype=torch.bool) for i, prompt in enumerate(text): threshold = Clip_filter.prompt_threshold[prompt] filt_mask = filt_mask | (logits[:,i] >= threshold) filt_ids = [] remain_ids = [] for i, id in enumerate(ids): if filt_mask[i]: filt_ids.append(id) else: remain_ids.append(id) return {"remain_ids":remain_ids, "filtered_ids":filt_ids} def clip_feature(self, dataloader): dataloader.dataset.get_img = True dataloader.dataset.get_cap = False clip_features = [] ids = [] for i, batch in enumerate(dataloader): images = batch["images"] features = self.clip_filter.get_image_features(images).cpu() clip_features.append(features) ids.extend(batch["ids"]) clip_features = torch.cat(clip_features) return {"clip_features":clip_features, "ids":ids} def clip_logit(self, dataloader): dataloader.dataset.get_img = True dataloader.dataset.get_cap = False clip_features = [] clip_logits = [] ids = [] for i, batch in enumerate(dataloader): images = batch["images"] # logits = self.clip_filter.get_logits(images) feature = self.clip_filter.get_image_features(images) logits = self.clip_logit_by_feat(feature)["clip_logits"] clip_features.append(feature) clip_logits.append(logits) ids.extend(batch["ids"]) clip_features = torch.cat(clip_features) clip_logits = torch.cat(clip_logits) return {"clip_features":clip_features, "clip_logits":clip_logits, "ids":ids, "text": self.clip_filter.text} def clip_logit_by_feat(self, feature): feature = feature.clone().to(self.clip_filter.device) feature = feature / feature.norm(p=2, dim=-1, keepdim=True) logit_scale = self.clip_filter.model.logit_scale.exp() logits = ((feature @ self.clip_filter.text_features.T) * logit_scale).cpu() return {"clip_logits":logits, "text": self.clip_filter.text} if __name__ == "__main__": import pickle with open("/vision-nfs/torralba/scratch/jomat/sam_dataset/filt_result/sa_000000/clip_logits_result.pickle","rb") as f: result=pickle.load(f) feat = result['clip_features'] logits =Art_filter().clip_logit_by_feat(feat) print(logits)