LAVT-RIS/data/dataset_refer_zom.py

import os
import sys
import json
import torch.utils.data as data
import torch
import itertools
import numpy as np
from PIL import Image
import pdb
import copy
from random import choice
from bert.tokenization_bert import BertTokenizer

from refer.refer_zom import ZREFER
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(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 = ZREFER(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(ref_ids)

        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.attention_masks = []
        self.tokenizer = BertTokenizer.from_pretrained(args.bert_tokenizer)

        self.ROOT = '/data2/dataset/RefCOCO/VRIS'
        self.metric_learning = args.metric_learning
        self.exclude_multiobj = args.exclude_multiobj
        self.metric_mode = args.metric_mode
        self.exclude_position = False

        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 = []
            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

                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)

                sentences_for_ref.append(torch.tensor(padded_input_ids).unsqueeze(0))
                attentions_for_ref.append(torch.tensor(attention_mask).unsqueeze(0))

            self.input_ids.extend(sentences_for_ref)
            self.attention_masks.extend(attentions_for_ref)


    def get_classes(self):
        return self.classes


    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 _get_hardpos_verb(self, ref, seg_id, sent_idx) :
        if seg_id in self.multi_obj_ref_ids:
            return ''
        
        # Extract metadata for hard positives if present            
        hardpos_dict = self.hardpos_meta.get(str(seg_id), {})
        if self.hp_selection == 'strict' :
            sent_id_list = list(hardpos_dict.keys())
            cur_hardpos = hardpos_dict.get(sent_id_list[sent_idx], {}).get('phrases', [])
        else : 
            cur_hardpos = list(itertools.chain.from_iterable(hardpos_dict[sid]['phrases'] for sid in hardpos_dict))

        if cur_hardpos:
            # Assign a hard positive verb phrase if available
            raw_verb = random.choice(cur_hardpos)
            return raw_verb
        
        return ''

    def __len__(self):
        return len(self.all_sent_id_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]]

        IMAGE_DIR = '/data2/dataset/COCO2014/trainval2014/'
        img = Image.open(os.path.join(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 = []
            att = []
            for s in range(len(self.input_ids[index])):
                padded_input_ids = self.input_ids[index][s]
                attention_mask = self.attention_masks[index][s]

                embedding.append(padded_input_ids.unsqueeze(-1))
                att.append(attention_mask.unsqueeze(-1))
            
            tensor_embeddings = torch.cat(embedding, dim=-1)
            attention_mask = torch.cat(att, dim=-1)
            return img, target, source_type, tensor_embeddings, attention_mask

        else:
            choice_sent = np.random.choice(len(self.input_ids[index]))
            tensor_embeddings = self.input_ids[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)

                ## Only the case with hardpos_ in metric_mode                
                if 'hardpos_' in self.metric_mode or self.hardneg_prob == 0.0:
                    pos_type = 'zero'
                    if 'refined' in self.metric_mode :
                        pos_sent_picked = self._get_hardpos_verb(ref, this_ref_id, choice_sent)
                    else :
                        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)  
                        pos_sent = pos_sent.squeeze(0) if pos_sent.dim() == 2 and pos_sent.size(0) == 1 else pos_sent
                        pos_attn_mask = pos_attn_mask.squeeze(0) if pos_attn_mask.size(0) == 1 else pos_attn_mask
                    
                    return img, target, source_type, tensor_embeddings, attention_mask, pos_sent, pos_attn_mask, pos_type
                                
            return img, target, source_type, tensor_embeddings, 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)