from typing import List, Dict, Union, Tuple
from PIL import Image, ImageDraw, ImageFilter, ImageOps, ImageEnhance
import spacy
import hashlib
import os
import torch
import torchvision
import torchvision.transforms as transforms
import clip
from transformers import BertTokenizer, RobertaTokenizerFast
import ruamel.yaml as yaml
import copy
from interpreter import Box
import pycocotools.mask as mask_utils
import alpha_clip
from segment_anything import sam_model_registry, SamPredictor
import numpy as np
import cv2
import matplotlib.pyplot as plt
import pickle
class Executor:
def __init__(self, device: str = "cpu", box_representation_method: str = "crop", method_aggregator: str = "max", enlarge_boxes: int = 0, expand_position_embedding: bool = False, square_size: bool = False, blur_std_dev: int = 100, cache_path: str = None, input_file: str = None) -> None:
IMPLEMENTED_METHODS = ["blur", "full", "gray"]
if any(m not in IMPLEMENTED_METHODS for m in box_representation_method.split(",")):
raise NotImplementedError
IMPLEMENTED_AGGREGATORS = ["max", "sum"]
if method_aggregator not in IMPLEMENTED_AGGREGATORS:
raise NotImplementedError
self.box_representation_method = box_representation_method
self.method_aggregator = method_aggregator
self.enlarge_boxes = enlarge_boxes
self.device = device
self.expand_position_embedding = expand_position_embedding
self.square_size = square_size
self.blur_std_dev = blur_std_dev
self.cache_path = cache_path
def preprocess_image(self, image: Image) -> List[torch.Tensor]:
return [preprocess(image) for preprocess in self.preprocesses]
def preprocess_mask(self, mask: Image) -> List[torch.Tensor]:
preprocess = self.preprocesses[0]
return preprocess.transforms[1](preprocess.transforms[0](mask))
def preprocess_text(self, text: str) -> torch.Tensor:
raise NotImplementedError
def call_model(self, model: torch.nn.Module, images: torch.Tensor, text: Union[torch.Tensor, Dict[str, torch.Tensor]]) -> torch.Tensor:
raise NotImplementedError
def tensorize_inputs(self, caption: str, image: Image, boxes: List[Box], image_name: str = None, image_pth: str = None) -> Tuple[List[torch.Tensor], torch.Tensor]:
images = []
for preprocess in self.preprocesses:
images.append([])
if 'aclip' in self.clip_type:
self.all_masks = []
read_save = False
if self.mask_path is not None: # load mask if cached
file_name = image_pth.split('/')[-1].split('.')[0]+'.pkl'
if os.path.exists(os.path.join(self.mask_path, file_name)):
all_rles = pickle.load(open(os.path.join(self.mask_path, file_name),'rb'))
for rle in all_rles:
mask = np.array(mask_utils.decode(rle), dtype=bool)
self.all_masks.append(mask)
read_save = True
if not read_save:
# use SAM to generate masks
self.predictor.set_image(np.array(image.convert('RGB')))
all_rles = []
for i in range(len(boxes)):
box = [
max(boxes[i].left-self.enlarge_boxes, 0),
max(boxes[i].top-self.enlarge_boxes, 0),
min(boxes[i].right+self.enlarge_boxes, image.width),
min(boxes[i].bottom+self.enlarge_boxes, image.height)
] # box prompt
input_box = np.array(box)
masks, _, _ = self.predictor.predict(
point_coords=None,
point_labels=None,
box=input_box[None, :],
multimask_output=False,
)
self.all_masks.append(masks[0])
rle = mask_utils.encode(np.array(masks[0][:, :, None], order='F', dtype="uint8"))[0]
rle["counts"] = rle["counts"].decode("utf-8")
all_rles.append(rle)
if self.mask_path is not None: # save mask
os.makedirs(self.mask_path, exist_ok=True)
pickle.dump(all_rles, open(os.path.join(self.mask_path, file_name),'wb'))
if self.cache_path is None or any([not os.path.exists(os.path.join(self.cache_path, "refcoco_val", model_name, "image", image_name, method_name+".pt")) for model_name in self.model_names for method_name in self.box_representation_method.split(',')]):
if "full" in self.box_representation_method: # original full image with alpha-map
for i in range(len(boxes)):
image_i = image.copy()
preprocessed_images = self.preprocess_image(image_i)
for j, img in enumerate(preprocessed_images):
images[j].append(img.to(self.device))
if "blur" in self.box_representation_method:
for i in range(len(boxes)):
image_i = image.copy()
mask = Image.new('L', image_i.size, 0)
draw = ImageDraw.Draw(mask)
box = (
max(boxes[i].left-self.enlarge_boxes, 0),
max(boxes[i].top-self.enlarge_boxes, 0),
min(boxes[i].right+self.enlarge_boxes, image_i.width),
min(boxes[i].bottom+self.enlarge_boxes, image_i.height)
)
if 'aclip' in self.clip_type:
width, height = image.size
for y in range(height):
for x in range(width):
if self.all_masks[i][y][x] == 1:
draw.point((x, y), fill=255)
else:
draw.rectangle([box[:2], box[2:]], fill=255)
blurred = image_i.filter(ImageFilter.GaussianBlur(self.blur_std_dev))
blurred.paste(image_i, mask=mask)
preprocessed_images = self.preprocess_image(blurred)
for j, img in enumerate(preprocessed_images):
images[j].append(img.to(self.device))
if "gray" in self.box_representation_method:
for i in range(len(boxes)):
image_i = image.copy()
mask_i = self.all_masks[i]
width, height = image.size
pixels = image_i.load()
for y in range(height):
for x in range(width):
if mask_i[y][x] == 0:
pixel_value = pixels[x, y]
gray_value = int(0.2989 * pixel_value[0] + 0.5870 * pixel_value[1] + 0.1140 * pixel_value[2])
pixels[x, y] = (gray_value, gray_value, gray_value)
preprocessed_images = self.preprocess_image(image_i)
for j, img in enumerate(preprocessed_images):
images[j].append(img.to(self.device))
imgs = [torch.stack(image_list) for image_list in images]
else:
imgs = [[] for _ in self.models]
text_tensor = self.preprocess_text(caption.lower()).to(self.device)
return imgs, text_tensor
@torch.no_grad()
def __call__(self, caption: str, image: Image, boxes: List[Box], image_name: str = None, image_pth=None) -> torch.Tensor:
images, text_tensor = self.tensorize_inputs(caption, image, boxes, image_name, image_pth)
all_logits_per_image = []
all_logits_per_text = []
box_representation_methods = self.box_representation_method.split(',')
caption_hash = hashlib.md5(caption.encode('utf-8')).hexdigest()
for model, images_t, model_name in zip(self.models, images, self.model_names):
self.image_feat_path = ""
if self.cache_path is not None:
text_cache_path = os.path.join(self.cache_path, "refcoco_val", model_name, "text"+("_shade" if self.box_representation_method == "shade" else ""))
image_feat_path = os.path.join(self.cache_path, "refcoco_val", model_name, "image", image_name)
self.image_feat_path = image_feat_path
image_features = None
text_features = None
if self.cache_path is not None and os.path.exists(os.path.join(self.cache_path, "refcoco_val", model_name)):
if os.path.exists(os.path.join(text_cache_path, caption_hash+".pt")):
text_features = torch.load(os.path.join(text_cache_path, caption_hash+".pt"), map_location=self.device)
if os.path.exists(image_feat_path):
if all([os.path.exists(os.path.join(image_feat_path, method_name+".pt")) for method_name in box_representation_methods]):
image_features = []
for method_name in box_representation_methods:
features = torch.load(os.path.join(image_feat_path, method_name+".pt"), map_location=self.device)
image_features.append(torch.stack([
features[(box.x, box.y, box.w, box.h)]
for box in boxes
]))
image_features = torch.stack(image_features)
image_features = image_features.view(-1, image_features.shape[-1])
logits_per_image, logits_per_text, image_features, text_features = self.call_model(model, images_t, text_tensor, image_features=image_features, text_features=text_features, boxes=boxes, image_pth=image_pth)
all_logits_per_image.append(logits_per_image)
all_logits_per_text.append(logits_per_text)
if self.cache_path is not None and image_name is not None and image_features is not None:
image_features = image_features.view(len(box_representation_methods), len(boxes), image_features.shape[-1])
if not os.path.exists(image_feat_path):
os.makedirs(image_feat_path)
for i in range(image_features.shape[0]):
method_name = box_representation_methods[i]
if not os.path.exists(os.path.join(image_feat_path, method_name+".pt")):
image_features_dict = {(box.x, box.y, box.w, box.h): image_features[i,j,:].cpu() for j, box in enumerate(boxes)}
torch.save(image_features_dict, os.path.join(image_feat_path, method_name+".pt"))
if self.cache_path is not None and not os.path.exists(os.path.join(text_cache_path, caption_hash+".pt")) and text_features is not None:
assert text_features.shape[0] == 1
if not os.path.exists(text_cache_path):
os.makedirs(text_cache_path)
torch.save(text_features.cpu(), os.path.join(text_cache_path, caption_hash+".pt"))
all_logits_per_image = torch.stack(all_logits_per_image).sum(0)
all_logits_per_text = torch.stack(all_logits_per_text).sum(0)
if self.method_aggregator == "max":
all_logits_per_text = all_logits_per_text.view(-1, len(boxes)).max(dim=0, keepdim=True)[0]
elif self.method_aggregator == "sum":
all_logits_per_text = all_logits_per_text.view(-1, len(boxes)).sum(dim=0, keepdim=True)
return all_logits_per_text.view(-1)
class ClipExecutor(Executor):
def __init__(self, clip_model: str = "ViT-B/32", device: str = "cpu", box_representation_method: str = "crop", method_aggregator: str = "max", enlarge_boxes: int = 0, expand_position_embedding: bool = False, square_size: bool = False, blur_std_dev: int = 100, cache_path: str = None, input_file: str = None, clip_type: str=None) -> None:
super().__init__(device, box_representation_method, method_aggregator, enlarge_boxes, expand_position_embedding, square_size, blur_std_dev, cache_path)
self.clip_models = clip_model.split(",")
self.model_names = [model_name.replace("/", "_") for model_name in self.clip_models]
self.models = []
self.preprocesses = []
self.data_name = input_file.split('/')[-1].split('.')[0]
self.mask_path = None
self.clip_type = clip_type
if self.cache_path is not None:
self.mask_path = os.path.join(self.cache_path, "refcoco_val", 'det_masks')
sam_checkpoint = "./ckpt/sam_vit_h_4b8939.pth"
model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
self.predictor = SamPredictor(sam)
for model_name in self.clip_models:
if 'aclip' in self.clip_type:#using alpha-clip
self.mask_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Resize((224, 224)),
transforms.Normalize(0.5, 0.26)
])
if model_name == 'ViT-B/16':
model, preprocess = alpha_clip.load("ViT-B/16", alpha_vision_ckpt_pth="./ckpt/grit1m/clip_b16_grit+mim_fultune_4xe.pth", device=device)
elif model_name == 'ViT-L/14':
model, preprocess = alpha_clip.load("ViT-L/14", alpha_vision_ckpt_pth="./ckpt/grit1m/clip_l14_grit+mim_fultune_6xe.pth", device=device)
else: model, preprocess = clip.load(model_name, device=device, jit=False)
self.models.append(model)
if self.square_size:
print("Square size!")
preprocess.transforms[0] = transforms.Resize((model.visual.input_resolution, model.visual.input_resolution), interpolation=transforms.InterpolationMode.BICUBIC)
self.preprocesses.append(preprocess)
self.models = torch.nn.ModuleList(self.models)
def preprocess_text(self, text: str) -> torch.Tensor:
if "aclip" in self.box_representation_method:
return alpha_clip.tokenize([text.lower()])
if "shade" in self.box_representation_method:
return clip.tokenize([text.lower()+" is in red color."])
return clip.tokenize(["a photo of "+text.lower()])
def call_model(self, model: torch.nn.Module, images: torch.Tensor, text: torch.Tensor, image_features: torch.Tensor = None, text_features: torch.Tensor = None, boxes=None, image_pth=None) -> torch.Tensor:
if image_features is None:
print('computing image features')
if 'aclip' not in self.clip_type:
image_features = model.encode_image(images)
image_features = image_features / image_features.norm(dim=-1, keepdim=True)
else:
image_features = []
if 'full' in self.box_representation_method:
aclip_images = images[:len(boxes)]
alphas = []
if os.path.exists(os.path.join(self.image_feat_path, 'full.pt')):
features = torch.load(os.path.join(self.image_feat_path, 'full.pt'), map_location=self.device)
aclip_image_features = torch.stack([
features[(box.x, box.y, box.w, box.h)]
for box in boxes
])
else:
for i in range(len(self.all_masks)):
binary_mask = self.all_masks[i]
alpha = self.mask_transform((binary_mask * 255).astype(np.uint8))
alpha = alpha.half().cuda().unsqueeze(dim=0)
alphas.append(alpha)
alphas = torch.cat(alphas, dim=0)
aclip_images = aclip_images.half()
aclip_image_features = model.visual(aclip_images, alphas) # using alpha channels
images = images[len(boxes):]
image_features.append(aclip_image_features)
if 'blur' in self.box_representation_method:
if os.path.exists(os.path.join(self.image_feat_path, 'blur.pt')):
features = torch.load(os.path.join(self.image_feat_path, 'blur.pt'), map_location=self.device)
ablur_images_features = torch.stack([
features[(box.x, box.y, box.w, box.h)]
for box in boxes
])
else:
ablur_images = images[:len(boxes)]
alphas = []
for i in range(len(self.all_masks)):
binary_mask = self.all_masks[i]
alpha = self.mask_transform((binary_mask * 255).astype(np.uint8))
alpha = alpha.half().cuda().unsqueeze(dim=0)
alphas.append(alpha)
alphas = torch.cat(alphas, dim=0)
ablur_images = ablur_images.half()
ablur_images_features = model.visual(ablur_images, alphas)
images = images[len(boxes):]
image_features.append(ablur_images_features)
if 'gray' in self.box_representation_method:
if os.path.exists(os.path.join(self.image_feat_path, 'gray.pt')):
features = torch.load(os.path.join(self.image_feat_path, 'gray.pt'), map_location=self.device)
gray_images_features = torch.stack([
features[(box.x, box.y, box.w, box.h)]
for box in boxes
])
else:
gray_images = images[:len(boxes)]
alphas = []
for i in range(len(self.all_masks)):
binary_mask = self.all_masks[i]
alpha = self.mask_transform((binary_mask * 255).astype(np.uint8))
alpha = alpha.half().cuda().unsqueeze(dim=0)
alphas.append(alpha)
alphas = torch.cat(alphas, dim=0)
gray_images = gray_images.half()
gray_images_features = model.visual(gray_images, alphas)
images = images[len(boxes):]
image_features.append(gray_images_features)
image_features = torch.cat(image_features, dim=0)
image_features = image_features / image_features.norm(dim=-1, keepdim=True)
if text_features is None:
print('computing text features')
text_features = model.encode_text(text)
# normalized features
text_features = text_features / text_features.norm(dim=-1, keepdim=True)
# cosine similarity as logits
logit_scale = model.logit_scale.exp()
logits_per_image = logit_scale * image_features @ text_features.t()
logits_per_text = logits_per_image.t()
return logits_per_image, logits_per_text, image_features, text_features
def __call__(self, caption: str, image: Image, boxes: List[Box], image_name: str = None, image_pth=None) -> torch.Tensor:
if self.expand_position_embedding:
original_preprocesses = self.preprocesses
new_preprocesses = []
original_position_embeddings = []
for model_name, model, preprocess in zip(self.clip_models, self.models, self.preprocesses):
if "RN" in model_name:
model_spatial_dim = int((model.visual.attnpool.positional_embedding.shape[0]-1)**0.5)
patch_size = model.visual.input_resolution // model_spatial_dim
original_positional_embedding = model.visual.attnpool.positional_embedding.clone()
model.visual.attnpool.positional_embedding = torch.nn.Parameter(torch.nn.functional.interpolate(
model.visual.attnpool.positional_embedding[1:,:].permute(1, 0).view(1, -1, model_spatial_dim, model_spatial_dim),
size=(image.height // patch_size, image.width // patch_size),
mode='bicubic',
align_corners=False
).squeeze(0).permute(1, 2, 0).view(-1, original_positional_embedding.shape[-1]))
model.visual.attnpool.positional_embedding = torch.nn.Parameter(torch.cat((
original_positional_embedding[:1,:],
model.visual.attnpool.positional_embedding
), dim=0))
transform = transforms.Compose([
transforms.Resize(((image.height // patch_size)*patch_size, (image.width // patch_size)*patch_size), interpolation=Image.BICUBIC),
lambda image: image.convert("RGB"),
transforms.ToTensor(),
transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
])
else:
model_spatial_dim = int((model.visual.positional_embedding.shape[0]-1)**0.5)
patch_size = model.visual.input_resolution // model_spatial_dim
original_positional_embedding = model.visual.positional_embedding.clone()
model.visual.positional_embedding = torch.nn.Parameter(torch.nn.functional.interpolate(
model.visual.positional_embedding[1:,:].permute(1, 0).view(1, -1, model_spatial_dim, model_spatial_dim),
size=(image.height // patch_size, image.width // patch_size),
mode='bicubic',
align_corners=False
).squeeze(0).permute(1, 2, 0).view(-1, original_positional_embedding.shape[-1]))
model.visual.positional_embedding = torch.nn.Parameter(torch.cat((
original_positional_embedding[:1,:],
model.visual.positional_embedding
), dim=0))
transform = transforms.Compose([
transforms.Resize(((image.height // patch_size)*patch_size, (image.width // patch_size)*patch_size), interpolation=Image.BICUBIC),
lambda image: image.convert("RGB"),
transforms.ToTensor(),
transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
])
new_preprocesses.append(transform)
original_position_embeddings.append(original_positional_embedding)
self.preprocesses = new_preprocesses
result = super().__call__(caption, image, boxes, image_name, image_pth)
if self.expand_position_embedding:
self.preprocesses = original_preprocesses
for model, model_name, pos_embedding in zip(self.models, self.clip_models, original_position_embeddings):
if "RN" in model_name:
model.visual.attnpool.positional_embedding = torch.nn.Parameter(pos_embedding)
else:
model.visual.positional_embedding = torch.nn.Parameter(pos_embedding)
return result