import os import torch.utils.data as data import torch import numpy as np from PIL import Image import pdb import copy from random import choice from bert.tokenization_bert import BertTokenizer from textblob import TextBlob from refer.refer import REFER import copy import random import torch from collections import defaultdict import torch import torch.distributed as dist from torch.utils.data.distributed import DistributedSampler from args import get_parser import random # Dataset configuration initialization parser = get_parser() args = parser.parse_args() class Referzom_Dataset_HP(data.Dataset): def __init__(self, args, image_transforms=None, target_transforms=None, split='train', eval_mode=False): self.classes = [] self.image_transforms = image_transforms self.target_transform = target_transforms self.split = split self.refer = REFER(args.refer_data_root, args.dataset, args.splitBy) self.dataset_type = args.dataset self.max_tokens = 20 ref_ids = self.refer.getRefIds(split=self.split) self.img_ids = self.refer.getImgIds() all_imgs = self.refer.Imgs self.imgs = list(all_imgs[i] for i in self.img_ids) self.ref_ids = ref_ids self.input_ids = [] self.input_ids_masked = [] self.attention_masks = [] self.tokenizer = BertTokenizer.from_pretrained(args.bert_tokenizer) # for metric learning ##################### self.ROOT = '/data2/dataset/RefCOCO/VRIS' # self.ROOT = '/data2/projects/seunghoon/VerbRIS/VerbCentric_CY/datasets/VRIS' self.metric_learning = args.metric_learning self.exclude_multiobj = args.exclude_multiobj self.metric_mode = args.metric_mode self.exclude_position = False self.hp_selection = args.hp_selection if self.metric_learning and eval_mode == False: self.hardneg_prob = args.hn_prob self.multi_obj_ref_ids = self._load_multi_obj_ref_ids() self.hardpos_meta, self.hardneg_meta = self._load_metadata() else: self.hardneg_prob = 0.0 self.multi_obj_ref_ids = None self.hardpos_meta, self.hardneg_meta = None, None ############################################# self.eval_mode = eval_mode self.zero_sent_id_list = [] self.one_sent_id_list = [] self.all_sent_id_list = [] self.sent_2_refid = {} for r in ref_ids: ref = self.refer.loadRefs(r) source_type = ref[0]['source'] for sent_dict in ref[0]['sentences']: sent_id = sent_dict['sent_id'] self.sent_2_refid[sent_id] = r self.all_sent_id_list.append(sent_id) if source_type=='zero': self.zero_sent_id_list.append(sent_id) else: self.one_sent_id_list.append(sent_id) for r in ref_ids: ref = self.refer.Refs[r] sentences_for_ref = [] sentences_for_ref_masked = [] attentions_for_ref = [] for i, el in enumerate(ref['sentences']): sentence_raw = el['raw'] attention_mask = [0] * self.max_tokens padded_input_ids = [0] * self.max_tokens padded_input_ids_masked = [0] * self.max_tokens blob = TextBlob(sentence_raw.lower()) chara_list = blob.tags mask_ops = [] mask_ops1 = [] for word_i, (word_now, chara) in enumerate(chara_list): if (chara == 'NN' or chara == 'NNS') and word_i < 19 and word_now.lower(): mask_ops.append(word_i) mask_ops1.append(word_now) mask_ops2 = self.get_adjacent_word(mask_ops) input_ids = self.tokenizer.encode(text=sentence_raw, add_special_tokens=True) # truncation of tokens input_ids = input_ids[:self.max_tokens] padded_input_ids[:len(input_ids)] = input_ids attention_mask[:len(input_ids)] = [1]*len(input_ids) if len(mask_ops) == 0: attention_remask = attention_mask input_ids_masked = input_ids else: could_mask = choice(mask_ops2) input_ids_masked = copy.deepcopy(input_ids) for i in could_mask: input_ids_masked[i + 1] = 0 padded_input_ids_masked[:len(input_ids_masked)] = input_ids_masked sentences_for_ref.append(torch.tensor(padded_input_ids).unsqueeze(0)) sentences_for_ref_masked.append(torch.tensor(padded_input_ids_masked).unsqueeze(0)) attentions_for_ref.append(torch.tensor(attention_mask).unsqueeze(0)) self.input_ids.extend(sentences_for_ref) self.input_ids_masked.extend(sentences_for_ref_masked) self.attention_masks.extend(attentions_for_ref) def get_classes(self): return self.classes def __len__(self): return len(self.all_sent_id_list) def get_adjacent_word(self, mask_list): output_mask_list = [] length = len(mask_list) i = 0 while i < length: begin_pos = i while i+1 < length and mask_list[i+1] == mask_list[i] + 1: i += 1 end_pos = i+1 output_mask_list.append(mask_list[begin_pos:end_pos]) i = end_pos return output_mask_list # for metric learning ##################### ########################################### def _tokenize(self, sentence): attention_mask = [0] * self.max_tokens padded_input_ids = [0] * self.max_tokens input_ids = self.tokenizer.encode(text=sentence, add_special_tokens=True) # truncation of tokens input_ids = input_ids[:self.max_tokens] padded_input_ids[:len(input_ids)] = input_ids attention_mask[:len(input_ids)] = [1]*len(input_ids) # match shape as (1, max_tokens) return torch.tensor(padded_input_ids).unsqueeze(0), torch.tensor(attention_mask).unsqueeze(0) def _load_multi_obj_ref_ids(self): # Load multi-object reference IDs based on configurations if not self.exclude_multiobj and not self.exclude_position : return None elif self.exclude_position: multiobj_path = os.path.join(self.ROOT, 'multiobj_ov2_nopos.txt') elif self.exclude_multiobj : multiobj_path = os.path.join(self.ROOT, 'multiobj_ov3.txt') with open(multiobj_path, 'r') as f: return [int(line.strip()) for line in f.readlines()] def _load_metadata(self): hardpos_path = os.path.join(self.ROOT, 'verb_ext_text_example_refzom.json') with open(hardpos_path, 'r', encoding='utf-8') as f: hardpos_json = json.load(f) if "hardpos_only" in self.metric_mode : hardneg_json = None # else : # hardneg_path = os.path.join(self.ROOT, 'hardneg_verb.json') # with open(hardneg_path, 'r', encoding='utf-8') as q: # hardneg_json = json.load(q) return hardpos_json, hardneg_json def __getitem__(self, index): sent_id = self.all_sent_id_list[index] this_ref_id = self.sent_2_refid[sent_id] this_img_id = self.refer.getImgIds(this_ref_id) this_img = self.refer.Imgs[this_img_id[0]] img = Image.open(os.path.join(self.refer.IMAGE_DIR, this_img['file_name'])).convert("RGB") ref = self.refer.loadRefs(this_ref_id) if self.dataset_type == 'ref-zom': source_type = ref[0]['source'] else: source_type = 'not_zero' ref_mask = np.array(self.refer.getMask(ref[0])['mask']) annot = np.zeros(ref_mask.shape) annot[ref_mask == 1] = 1 annot = Image.fromarray(annot.astype(np.uint8), mode="P") if self.image_transforms is not None: img, target = self.image_transforms(img, annot) if self.eval_mode: embedding = [] embedding_masked = [] att = [] for s in range(len(self.input_ids[index])): e = self.input_ids[index][s] a = self.attention_masks[index][s] embedding.append(e.unsqueeze(-1)) embedding_masked.append(e.unsqueeze(-1)) att.append(a.unsqueeze(-1)) tensor_embeddings = torch.cat(embedding, dim=-1) tensor_embeddings_masked = torch.cat(embedding_masked, dim=-1) attention_mask = torch.cat(att, dim=-1) return img, target, source_type, tensor_embeddings, tensor_embeddings_masked, attention_mask else: choice_sent = np.random.choice(len(self.input_ids[index])) tensor_embeddings = self.input_ids[index][choice_sent] tensor_embeddings_masked = self.input_ids_masked[index][choice_sent] attention_mask = self.attention_masks[index][choice_sent] if self.metric_learning : pos_sent = torch.zeros_like(tensor_embeddings) pos_attn_mask = torch.zeros_like(attention_mask) pos_type = 'zero' if 'hardpos_' in self.metric_mode or self.hardneg_prob == 0.0 : pos_sents = self.hardpos_meta[str(this_ref_id)].values() # drop elements with none pos_sents = [s for s in pos_sents if s is not None] pos_sent_picked = random.choice(list(pos_sents)) if pos_sent_picked: pos_type = 'hardpos' pos_sent, pos_attn_mask = self._tokenize(pos_sent_picked) return img, target, source_type, tensor_embeddings, tensor_embeddings_masked, attention_mask, pos_sent, pos_attn_mask, pos_type return img, target, source_type, tensor_embeddings, tensor_embeddings_masked, attention_mask class Refzom_DistributedSampler(DistributedSampler): def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True): super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle) self.one_id_list = dataset.one_sent_id_list self.zero_id_list = dataset.zero_sent_id_list self.sent_ids_list = dataset.all_sent_id_list if self.shuffle==True: random.shuffle(self.one_id_list) random.shuffle(self.zero_id_list) self.sent_id = self.insert_evenly(self.zero_id_list,self.one_id_list) self.indices = self.get_positions(self.sent_ids_list, self.sent_id) def get_positions(self, list_a, list_b): position_dict = {value: index for index, value in enumerate(list_a)} positions = [position_dict[item] for item in list_b] return positions def insert_evenly(self, list_a, list_b): len_a = len(list_a) len_b = len(list_b) block_size = len_b // len_a result = [] for i in range(len_a): start = i * block_size end = (i + 1) * block_size result.extend(list_b[start:end]) result.append(list_a[i]) remaining = list_b[(len_a * block_size):] result.extend(remaining) return result def __iter__(self): indices_per_process = self.indices[self.rank::self.num_replicas] return iter(indices_per_process)