import os import copy import pickle from PIL import Image from torchvision import transforms import random from torchvision.transforms.transforms import ToTensor from tqdm import tqdm import numpy as np from uniperceiver.config import configurable from uniperceiver.functional import read_np, dict_as_tensor, boxes_to_locfeats from ..build import DATASETS_REGISTRY import glob from uniperceiver.tokenization import ClipTokenizer import json from collections import defaultdict from uniperceiver.datasets.custom_transforms import clip_transforms import pyarrow as pa from uniperceiver.utils import comm __all__ = ["ImageTextPairDataset"] memorycache = False @DATASETS_REGISTRY.register() class ImageTextPairDataset: @configurable def __init__( self, cfg: str, stage: str, anno_file: str, seq_per_img: int, max_seq_len: int, feats_folder: str, relation_file: str, gv_feat_file: str, attribute_file: str, transform, tokenizer, data_percentage, tokenizer_name, use_ceph: bool, tcs_conf_path, task_type, preload_feats = None, random_mask=False, text_type_id=0, ): assert len(task_type)>0 self.cfg = cfg self.stage = stage self.anno_file = anno_file self.seq_per_img = seq_per_img assert self.seq_per_img == 1 self.use_ceph = use_ceph self.task_type = task_type if self.use_ceph: self.feats_folder = 's3://coco' print('debug info for coco pretrain: {} '.format(self.feats_folder)) from uniperceiver.datasets import TCSLoader if 'SLURM_PROCID' in os.environ: self.tcs_loader = TCSLoader(tcs_conf_path) else: self.tcs_loader = TCSLoader('petreloss_local.config') else: # local image folder self.feats_folder = feats_folder self.max_seq_len = max_seq_len self.relation_file = relation_file self.gv_feat_file = gv_feat_file self.attribute_file = attribute_file self.data_percentage = data_percentage self.tokenizer = tokenizer self.tokenizer_name = tokenizer_name self.use_clip_tokenizer = tokenizer_name == 'clip' self.initialized = False self.transform = transform self.loaded_feats = None if preload_feats: self.loaded_feats = self.pre_load_feats(preload_feats) # for index_maping self.idx2name = dict() self.name2idx = dict() # please if isinstance(self.anno_file, list): imageinfo = list() for anno_file in self.anno_file: imageinfo.extend(json.load(open(anno_file))["images"]) else: imageinfo = json.load(open(self.anno_file))["images"] for info in imageinfo: self.idx2name[info['id']] = { "split": info['file_path'], "name": info['file_name']} self.name2idx[info['file_name']] = info['id'] self.random_mask = random_mask self.text_type_id = text_type_id self.task_info = { 'task_type' : self.cfg.DATASETS.TASK_TYPE, 'dataset_name' : self.cfg.DATASETS.DATASET_NAME, 'batch_size' : self.cfg.DATALOADER.TRAIN_BATCH_SIZE if self.stage == 'train' else self.cfg.DATALOADER.TEST_BATCH_SIZE, 'sampling_weight': self.cfg.DATALOADER.SAMPLING_WEIGHT } _temp_list =self.load_data(self.cfg) self.database = pa.array(_temp_list) if comm.is_main_process(): import sys print("MSCOCO Pretrain Dataset:") print('!!! length of _temp_list: ', len(_temp_list)) print('!!! size of _temp_list: ', sys.getsizeof(_temp_list)) print('!!! size of pa database: ', sys.getsizeof(self.database)) del _temp_list def pre_load_feats(self, preload_feat_folder): loaded_feats = {} file_list = glob.glob(os.path.join(preload_feat_folder, '*.pkl')) for fname in file_list: with open(fname, 'rb') as f: feats = pickle.load(f) loaded_feats.update(feats) return loaded_feats @classmethod def from_config(cls, cfg, stage: str = "train"): if 'SLURM_PROCID' in os.environ: tcs_conf_path = cfg.DATALOADER.get("TCS_CONF_PATH", "petreloss.config") else: # dev machine tcs_conf_path = "slurm_tools/petreloss_local.config" ann_files = { "train": [os.path.join(cfg.DATALOADER.ANNO_FOLDER, "captions_train113k.json"), os.path.join(cfg.DATALOADER.ANNO_FOLDER, "captions_val5k.json")], # no validation "test": os.path.join(cfg.DATALOADER.ANNO_FOLDER, "captions_test5k.json") } if getattr(cfg.DATALOADER, 'TRANSFORM', None) == 'clip_transforms': transform = clip_transforms(stage, img_size=cfg.MODEL.IMG_INPUT_SIZE) else: transform = transforms.Compose([ transforms.Resize([224, 224]), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))] ) ret = { "cfg" : cfg, "stage" : stage, "anno_file" : ann_files[stage], "seq_per_img" : 1, "feats_folder" : cfg.DATALOADER.FEATS_FOLDER, "relation_file" : cfg.DATALOADER.RELATION_FILE, "gv_feat_file" : cfg.DATALOADER.GV_FEAT_FILE, "attribute_file" : cfg.DATALOADER.ATTRIBUTE_FILE, "max_seq_len" : cfg.MODEL.MAX_SEQ_LEN, "use_ceph" : getattr(cfg.DATALOADER, 'USE_CEPH', False), "tcs_conf_path" : tcs_conf_path, "transform" : transform, 'task_type' : cfg.DATASETS.TASK_TYPE, "random_mask" : getattr(cfg.DATALOADER, 'RANDOM_MASK', False), "data_percentage": cfg.DATALOADER.DATA_PERCENTAGE, "text_type_id" : getattr(cfg.DATALOADER, 'TYPE_EMBEDDING_ID', 0), } ret['tokenizer'] = ClipTokenizer() ret['tokenizer_name'] = "clip" return ret def _preprocess_datalist(self, datalist): return datalist def load_data(self, cfg): if self.stage == "test": cache_path = os.path.join( os.path.dirname(self.anno_file), "cache", "mscoco_caption_w_testcap_%s.pkl" % ( self.stage) ) if not os.path.exists(os.path.dirname(cache_path)): os.makedirs(os.path.dirname(cache_path)) if not os.path.exists(cache_path): datalist = self.load_raw_data(cfg, self.anno_file) pickle.dump(datalist, open(cache_path, "wb")) datalist = pickle.load(open(cache_path, "rb")) else: datalist = list() assert self.stage == "train", "no validation now" for i, stage in enumerate(["train", "val"]): cache_path = os.path.join( os.path.dirname(self.anno_file[i]), "cache", "mscoco_caption_w_testcap_%s.pkl" % ( stage) ) if not os.path.exists(os.path.dirname(cache_path)): os.makedirs(os.path.dirname(cache_path)) if not os.path.exists(cache_path): datalist_part = self.load_raw_data(cfg, self.anno_file[i]) pickle.dump(datalist_part, open(cache_path, "wb")) datalist_part = pickle.load(open(cache_path, "rb")) datalist.extend(datalist_part) def _load_pkl_file(filepath): return pickle.load(open(filepath, 'rb'), encoding='bytes') if len(filepath) > 0 else None ext_data = { "relation": _load_pkl_file(self.relation_file), "attribute": _load_pkl_file(self.attribute_file), "gv_feat": _load_pkl_file(self.gv_feat_file) } for i in range(len(datalist)): image_id = int(datalist[i]['image_id']) for data_type in ext_data: if ext_data[data_type] is not None: if str(image_id) in ext_data[data_type]: datalist[i][data_type] = ext_data[data_type][str(image_id)] elif image_id in ext_data[data_type]: datalist[i][data_type] = ext_data[data_type][image_id] if self.data_percentage < 1.0 and self.stage == 'train': datalist = random.sample(datalist, k = int(self.data_percentage* len(datalist) ) ) return datalist def load_raw_data(self, cfg, anno_file): datalist = [] annoinfo = json.load(open(anno_file)) captions_train = sorted( annoinfo['annotations'], key=lambda x: x['id']) image_caption_info = defaultdict(list) for cap_info in captions_train: image_caption_info[cap_info['image_id']].append(cap_info['caption']) for im_id, caps in image_caption_info.items(): datalist.append( { "image_id": im_id, "captions": caps, } ) return datalist def random_word_wwm(self, tokens): output_tokens = [] output_label = [] for i, token in enumerate(tokens): if self.use_clip_tokenizer: sub_tokens = self.tokenizer.encode_basic_tokenized_token(token) else: sub_tokens = self.tokenizer.wordpiece_tokenizer.tokenize(token) prob = random.random() # mask token with 15% probability if prob < 0.15: prob /= 0.15 # 80% randomly change token to mask token if prob < 0.8: for sub_token in sub_tokens: if self.use_clip_tokenizer: output_tokens.append(self.tokenizer.encoder["<|spe|>"]) else: output_tokens.append("[MASK]") # 10% randomly change token to random token elif prob < 0.9: for sub_token in sub_tokens: if self.use_clip_tokenizer: output_tokens.append(random.choice(list(range(len(self.tokenizer.encoder))))) else: output_tokens.append(random.choice(list(self.tokenizer.vocab.keys()))) # -> rest 10% randomly keep current token else: for sub_token in sub_tokens: output_tokens.append(sub_token) # append current token to output (we will predict these later) for sub_token in sub_tokens: if self.use_clip_tokenizer: output_label.append(sub_token) else: try: output_label.append(self.tokenizer.vocab[sub_token]) except KeyError: # For unknown words (should not occur with BPE vocab) output_label.append(self.tokenizer.vocab["[UNK]"]) else: for sub_token in sub_tokens: # no masking token (will be ignored by loss function later) output_tokens.append(sub_token) output_label.append(-1) return output_tokens, output_label def __len__(self): return len(self.database) # def __call__(self, dataset_dict): def __getitem__(self, index): for i_try in range(100): try: dataset_dict = self.database[index].as_py() image_id = dataset_dict['image_id'] image_split = self.idx2name[int(image_id)]['split'] image_name = self.idx2name[int(image_id)]['name'] # load image image_path = os.path.join(self.feats_folder, image_split, image_name) if self.use_ceph: img = self.tcs_loader(image_path).convert('RGB') else: img = Image.open(image_path).convert("RGB") break except Exception as e: print("Failed to load image from idb {} with error {} ; trial {};".format(self.database[index], e, i_try)) index = (index + 1) % len(self.database) while (self.database[index].as_py() is None): index = (index + 1) % len(self.database) continue img = self.transform(img) ret = { 'input_sample': [{ 'data' : img, 'invalid_mask': None, 'modality' : 'image', 'data_type': 'input', 'sample_info' :{'id': image_id, 'path': image_path} }] } self.target_set = self.cfg.DATASETS.TARGET_SET if self.task_type == 'image_caption' and self.stage != 'train': ret.update({ 'target_set': copy.deepcopy(self.target_set), 'target_sample': [], 'target_idx': [], 'task_info' : copy.deepcopy(self.task_info) }) dict_as_tensor(ret) return ret if self.task_type =='image_retrieval' and self.stage != 'train': captions = [caption + " <|endoftext|>" for caption in dataset_dict['captions']] caption_tokens_raw = [ self.tokenizer.encode(caption) for caption in captions] caption_tokens = [ caption_token[:(self.max_seq_len - 1)] + [caption_token[-1]] if len(caption_token) > self.max_seq_len else caption_token for caption_token in caption_tokens_raw ] else: caption = random.sample(dataset_dict['captions'], self.seq_per_img)[0] # caption = ['pilgrims', 'coffee', 'house', '-', 'outside', 'the', 'store'] caption = caption + " <|endoftext|>" if self.task_type == 'mlm': u_mask_type = 1 elif self.task_type == 'image_caption': u_mask_type = 0 # causal mask if self.task_type=='image_caption' or self.task_type =='mlm': if u_mask_type == 1: # mlm caption_tokens = self.tokenizer.basic_tokenize(caption) caption_tokens, mlm_labels = self.random_word_wwm(caption_tokens) else: # caption caption_tokens = self.tokenizer.encode(caption) mlm_labels = self.tokenizer.encode("<|spe|>")*len(caption_tokens) else: caption_tokens = self.tokenizer.encode(caption) if len(caption_tokens) > self.max_seq_len: # mlm task text_len_keep = self.max_seq_len caption_tokens = caption_tokens[:(text_len_keep - 1)] + [caption_tokens[-1]] if self.task_type=='image_caption' or self.task_type == 'mlm': mlm_labels = mlm_labels[:(text_len_keep - 1)] + [mlm_labels[-1]] # self.task_info = { # 'task_type' : self.cfg.DATASETS.TASK_TYPE, # 'dataset_name' : self.cfg.DATASETS.DATASET_NAME, # 'batch_size' : self.cfg.DATASETS.TRAIN_BATCH_SIZE # if self.stage == 'train' else self.cfg.DATASETS.TEST_BATCH_SIZE, # 'sampling_weight': self.cfg.DATALOADER.SAMPLING_WEIGHT # } if self.task_type == 'image_caption': source = np.array(caption_tokens, dtype=np.int64) source2 = np.array(mlm_labels, dtype=np.int64) ret['input_sample'].append({ 'data' :[source, source2], 'invalid_mask' : None, 'modality' : 'text', 'data_type' : 'input', 'sample_info' : { 'text_spe_cat': True, } }) ret.update({ "target_sample": [], "target_idx" : [np.array(caption_tokens, dtype=np.int64)], "target_set" : copy.deepcopy(self.target_set), 'task_info' : copy.deepcopy(self.task_info) }) elif self.task_type == 'mlm': ret['input_sample'].append({ 'data' : [np.array(caption_tokens, dtype=np.int64)], 'invalid_mask': None, 'modality' : 'text', 'data_type' : 'input', 'sample_info' : {"text_token_padding_length": self.max_seq_len} }) ret.update({ 'target_sample': [], "target_idx" : [np.array(mlm_labels, dtype=np.int64)], "target_set" : copy.deepcopy(self.target_set), 'task_info' : copy.deepcopy(self.task_info) }) elif self.task_type == 'image_retrieval': if self.stage == 'train': ret.update({ 'target_sample': [{ 'data' : [np.array(caption_tokens, dtype=np.int64)], 'modality' : 'text', 'invalid_mask': None, 'data_type' : 'target', 'sample_info' : {} }], 'target_idx' : [], 'target_set' : [], 'task_info' : copy.deepcopy(self.task_info) }) else: ret.update( { 'input_sample': [{ 'data' : img, 'invalid_mask': None, 'modality': 'image', 'data_type': 'input', 'sample_info' : { 'id' : (image_id, [image_id] * len(caption_tokens)), 'path' : image_path } }], 'target_sample': [{ 'data' : [np.array(caption_token, dtype=np.int64) for caption_token in caption_tokens], 'modality' : 'text', 'invalid_mask': None, 'data_type' : 'target', 'sample_info' : { 'sample_alone': True, } }], 'target_idx' : [], 'target_set' : [], 'task_info' : copy.deepcopy(self.task_info) } ) else: raise NotImplementedError dict_as_tensor(ret) return ret