from tkinter.messagebox import NO import torch import json from collections import defaultdict from PIL import Image, ImageDraw from copy import deepcopy import os import torchvision.transforms as transforms import torchvision from .base_dataset import BaseDataset, check_filenames_in_zipdata, recalculate_box_and_verify_if_valid from io import BytesIO import random from .tsv import TSVFile from io import BytesIO import base64 from PIL import Image import numpy as np def decode_base64_to_pillow(image_b64): return Image.open(BytesIO(base64.b64decode(image_b64))).convert('RGB') def decode_tensor_from_string(arr_str, use_tensor=True): arr = np.frombuffer(base64.b64decode(arr_str), dtype='float32') if use_tensor: arr = torch.from_numpy(arr) return arr def decode_item(item): item = json.loads(item) item['image'] = decode_base64_to_pillow(item['image']) for anno in item['annos']: anno['image_embedding_before'] = decode_tensor_from_string(anno['image_embedding_before']) anno['text_embedding_before'] = decode_tensor_from_string(anno['text_embedding_before']) anno['image_embedding_after'] = decode_tensor_from_string(anno['image_embedding_after']) anno['text_embedding_after'] = decode_tensor_from_string(anno['text_embedding_after']) return item def check_unique(images, fields): for field in fields: temp_list = [] for img_info in images: temp_list.append(img_info[field]) assert len(set(temp_list)) == len(temp_list), field def clean_data(data): for data_info in data: data_info.pop("original_img_id", None) data_info.pop("original_id", None) data_info.pop("sentence_id", None) # sentence id for each image (multiple sentences for one image) data_info.pop("dataset_name", None) data_info.pop("data_source", None) data_info["data_id"] = data_info.pop("id") def clean_annotations(annotations): for anno_info in annotations: anno_info.pop("iscrowd", None) # I have checked that all 0 for flickr, vg, coco anno_info.pop("category_id", None) # I have checked that all 1 for flickr vg. This is not always 1 for coco, but I do not think we need this annotation anno_info.pop("area", None) # anno_info.pop("id", None) anno_info["data_id"] = anno_info.pop("image_id") def draw_box(img, boxes): draw = ImageDraw.Draw(img) for box in boxes: draw.rectangle([box[0], box[1], box[2], box[3]], outline ="red", width=2) # x0 y0 x1 y1 return img def xyhw2xyxy(box): x0, y0, w, h = box return [ x0, y0, x0+w, y0+h ] def make_a_sentence(obj_names, clean=False): if clean: obj_names = [ name[:-6] if ("-other" in name) else name for name in obj_names] caption = "" tokens_positive = [] for obj_name in obj_names: start_len = len(caption) caption += obj_name end_len = len(caption) caption += ", " tokens_positive.append( [[start_len, end_len]] # in real caption, positive tokens can be disjoint, thus using list of list ) caption = caption[:-2] # remove last ", " return caption #, tokens_positive def mask_for_random_drop_text_or_image_feature(masks, random_drop_embedding): """ input masks tell how many valid grounding tokens for this image e.g., 1,1,1,1,0,0,0,0,0,0... If random_drop_embedding=both. we will random drop either image or text feature for each token, but we always make sure there is at least one feature used. In other words, the following masks are not valid (because for the second obj, no feature at all): image: 1,0,1,1,0,0,0,0,0 text: 1,0,0,0,0,0,0,0,0 if random_drop_embedding=image. we will random drop image feature and always keep the text one. """ N = masks.shape[0] if random_drop_embedding=='both': temp_mask = torch.ones(2,N) for i in range(N): if random.uniform(0, 1) < 0.5: # else keep both features idx = random.sample([0,1], 1)[0] # randomly choose to drop image or text feature temp_mask[idx,i] = 0 image_masks = temp_mask[0]*masks text_masks = temp_mask[1]*masks if random_drop_embedding=='image': image_masks = masks*(torch.rand(N)>0.5)*1 text_masks = masks return image_masks, text_masks def project(x, projection_matrix): """ x (Batch*768) should be the penultimate feature of CLIP (before projection) projection_matrix (768*768) is the CLIP projection matrix, which should be weight.data of Linear layer defined in CLIP (out_dim, in_dim), thus we need to apply transpose below. this function will return the CLIP feature (without normalziation) """ return x@torch.transpose(projection_matrix, 0, 1) def inv_project(y, projection_matrix): """ y (Batch*768) should be the CLIP feature (after projection) projection_matrix (768*768) is the CLIP projection matrix, which should be weight.data of Linear layer defined in CLIP (out_dim, in_dim). this function will return the CLIP penultimate feature. Note: to make sure getting the correct penultimate feature, the input y should not be normalized. If it is normalized, then the result will be scaled by CLIP feature norm, which is unknown. """ return y@torch.transpose(torch.linalg.inv(projection_matrix), 0, 1) class TSVDataset(BaseDataset): def __init__(self, tsv_path, which_embedder='clip', which_layer=['after','after'], # text and image prob_use_caption=1, random_drop_embedding='none', image_size=256, min_box_size=0.01, max_boxes_per_data=8, max_images=None, # set as 30K used to eval random_crop = False, random_flip = True, ): image_root = "a placeholder path as we are using tsv here" super().__init__(image_root, random_crop, random_flip, image_size) self.tsv_path = tsv_path self.which_embedder = which_embedder self.prob_use_caption = prob_use_caption self.random_drop_embedding = random_drop_embedding self.min_box_size = min_box_size self.max_boxes_per_data = max_boxes_per_data self.max_images = max_images assert which_layer in [ ['after','after'], ['before','after_renorm'], ['before','after_reproject'] ] assert random_drop_embedding in ['none', 'both', 'image'] self.which_layer_text = which_layer[0] self.which_layer_image = which_layer[1] #self.projection_matrix = torch.load(os.path.join(os.path.dirname(__file__), 'projection_matrix') ) self.projection_matrix = torch.load('projection_matrix.pth') # Load tsv data self.tsv_file = TSVFile(self.tsv_path) # Load preprocessed name embedding if which_embedder == 'bert': self.embedding_len = 1280 elif which_embedder == 'clip': self.embedding_len = 768 else: assert False def total_images(self): return len(self) def get_item_from_tsv(self, index): _, item = self.tsv_file[index] item = decode_item(item) return item def mapping(self, image_embedding): if self.which_layer_image == 'after': # both use CLIP aligned feature return image_embedding elif self.which_layer_image == 'after_renorm': # text use before, but image use after projection but normalize to 28.7 return image_embedding*28.7 elif self.which_layer_image == 'after_reproject': image_embedding = project( image_embedding.unsqueeze(0), self.projection_matrix.T ) image_embedding = image_embedding.squeeze(0) image_embedding = image_embedding / image_embedding.norm() image_embedding = image_embedding * 28.7 return image_embedding def __getitem__(self, index): if self.max_boxes_per_data > 99: assert False, "Are you sure setting such large number of boxes?" raw_item = self.get_item_from_tsv(index) is_det = raw_item.get('is_det', False) # if it is from detection (such as o365), then we will make a caption out = {} # -------------------- id and image ------------------- # out['id'] = raw_item['data_id'] image = raw_item['image'] image_tensor, trans_info = self.transform_image(image) out["image"] = image_tensor # -------------------- grounding token ------------------- # annos = raw_item['annos'] areas = [] all_boxes = [] all_masks = [] all_text_embeddings = [] all_image_embeddings = [] if is_det: all_category_names = [] text_embedding_name = 'text_embedding_before' if self.which_layer_text == 'before' else 'text_embedding_after' image_embedding_name = 'image_embedding_after' for anno in annos: x, y, w, h = anno['bbox'] valid, (x0, y0, x1, y1) = recalculate_box_and_verify_if_valid(x, y, w, h, trans_info, self.image_size, self.min_box_size) if valid: areas.append( (x1-x0)*(y1-y0) ) all_boxes.append( torch.tensor([x0,y0,x1,y1]) / self.image_size ) # scale to 0-1 all_masks.append(1) all_text_embeddings.append(anno[text_embedding_name]) all_image_embeddings.append( self.mapping(anno[image_embedding_name]) ) if is_det: all_category_names.append(anno["category_name"]) wanted_idxs = torch.tensor(areas).sort(descending=True)[1] wanted_idxs = wanted_idxs[0:self.max_boxes_per_data] boxes = torch.zeros(self.max_boxes_per_data, 4) masks = torch.zeros(self.max_boxes_per_data) text_embeddings = torch.zeros(self.max_boxes_per_data, self.embedding_len) image_embeddings = torch.zeros(self.max_boxes_per_data, self.embedding_len) if is_det: category_names = [] for i, idx in enumerate(wanted_idxs): boxes[i] = all_boxes[idx] masks[i] = all_masks[idx] text_embeddings[i] = all_text_embeddings[idx] image_embeddings[i] = all_image_embeddings[idx] if is_det: category_names.append(all_category_names[idx]) if self.random_drop_embedding != 'none': image_masks, text_masks = mask_for_random_drop_text_or_image_feature(masks, self.random_drop_embedding) else: image_masks = masks text_masks = masks out["boxes"] = boxes out["masks"] = masks out["image_masks"] = image_masks out["text_masks"] = text_masks out["text_embeddings"] = text_embeddings out["image_embeddings"] = image_embeddings # -------------------- caption ------------------- # if random.uniform(0, 1) < self.prob_use_caption: if is_det: out["caption"] = make_a_sentence(category_names) else: out["caption"] = raw_item["caption"] else: out["caption"] = "" return out def __len__(self): return len(self.tsv_file)