|
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 |
|
|
|
parser = get_parser() |
|
args = parser.parse_args() |
|
|
|
|
|
class Referzom_Dataset(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) |
|
|
|
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) |
|
|
|
|
|
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 |
|
|
|
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) |
|
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] |
|
|
|
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) |
