ASDA/dataset/data_loader_gref_sbert.py

# -*- coding: utf-8 -*-

"""
refcoco, refcoco+ and refcocog referring image detection and segmentation PyTorch dataset.
"""
import sys
import cv2
import os
import torch
import json
import random
import numpy as np
import os.path as osp
import torch.utils.data as data
sys.path.append('.')
import utils
import re

# from pytorch_pretrained_bert.tokenization import BertTokenizer
from utils.transforms import letterbox, random_affine, random_copy, random_crop, random_erase
import copy 

import clip

sys.modules['utils'] = utils
cv2.setNumThreads(0)

class ReferDataset(data.Dataset):
    SUPPORTED_DATASETS = {
        'refcoco': {
            'splits': ('train', 'val', 'testA', 'testB'),
            'params': {'dataset': 'refcoco', 'split_by': 'unc'}
        },
        'refcoco+': {
            'splits': ('train', 'val', 'testA', 'testB'),
            'params': {'dataset': 'refcoco+', 'split_by': 'unc'}
        },
        'refcocog': {
            'splits': ('train', 'val', 'test'),
            'params': {'dataset': 'refcocog', 'split_by': 'unc'}
        },
        'refcocog_g': {
            'splits': ('train', 'val'),
            'params': {'dataset': 'refcocog', 'split_by': 'google'}
        },
        'refcocog_u': {
            'splits': ('train', 'val', 'test'),
            'params': {'dataset': 'refcocog', 'split_by': 'unc'}
        },
        'grefcoco': {
            'splits': ('train', 'val', 'testA', 'testB'),
            'params': {'dataset': 'grefcoco', 'split_by': 'unc'}
        }
    }


    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):
        # Load metadata for hard positive verb phrases, hard negative queries
        # we set refined file as default option
        hardpos_path = '/data2/projects/seunghoon/VerbRIS/CrossVLT/hardpos_verdict_gref_v4.json'
        with open(hardpos_path, 'r', encoding='utf-8') as f:
            hardpos_json = json.load(f)
        return hardpos_json
    
    def __init__(self, data_root, split_root='data', dataset='refcoco', imsize=256, splitby='umd',
                 transform=None, augment=False, split='train', max_query_len=128, metric_learning=None):
        images_tmp = []
        self.data_root = data_root
        self.split_root = split_root
        self.dataset = dataset
        self.imsize = imsize
        self.query_len = max_query_len
        self.transform = transform
        self.word_len = 17
        self.emb_size = 384
        self.split = split
        self.augment=augment

        valid_splits = self.SUPPORTED_DATASETS[self.dataset]['splits']

        if split not in valid_splits:
            raise ValueError(
                'Dataset {0} does not have split {1}'.format(
                    self.dataset, split))
        
        self.anns_root = osp.join(self.data_root, 'anns', self.dataset, self.split+'.txt')
        if self.dataset == 'refcocog' :
            mask_anno_str = '{0}_{1}'.format(self.dataset, splitby)
            self.mask_root = osp.join(self.data_root, 'masks', mask_anno_str)
        else :
            self.mask_root = osp.join(self.data_root, 'masks', self.dataset)
            
        self.im_dir = osp.join(self.data_root, 'images', 'train2014')
        

        if self.dataset == 'refcocog' :
            dataset_path = osp.join(self.split_root, self.dataset + '_' + splitby)
            splits = [split]
            for split in splits:
                imgset_file = '{0}_{1}_{2}.pth'.format(self.dataset, splitby, split)
                imgset_path = osp.join(dataset_path, imgset_file)
                images_tmp += torch.load(imgset_path)
                
            # metric learning options
            self.ROOT = '/data2/projects/seunghoon/VerbRIS/VerbCentric_CY/'
            self.all_hp_root = "/data2/dataset/RefCOCO/refcocog/SBERT_gref_umd"
            # self.exclude_position = args.exclude_pos
            self.exclude_position = True
            self.exclude_multiobj = True
            self.metric_learning = metric_learning

            # self.metric_mode = args.metric_mode
            self.hp_selection = 'strict'

            # meta data loading
            if self.metric_learning and self.split == 'train':
                self.multi_obj_ref_ids = self._load_multi_obj_ref_ids()
                self.hardpos_meta = self._load_metadata()
                
                # make new self.images file with sentence idx and total sent num (per ref_id)
                from collections import defaultdict
                ref_sentence_counts = defaultdict(int)
                for item in images_tmp:
                    ref_sentence_counts[item[1]] += 1

                self.images = []
                ref_sentence_indices = defaultdict(int) 
                for item in images_tmp:
                    img_name, seg_id, box, sentence = item
                    sent_index = ref_sentence_indices[seg_id]  
                    total_sentences = ref_sentence_counts[seg_id] 
                    self.images.append((img_name, seg_id, box, sentence, sent_index, total_sentences))
                    ref_sentence_indices[seg_id] += 1

            else :
                self.images = images_tmp
                self.multi_obj_ref_ids = None
                self.hardpos_meta = None

        else : 
            dataset_path = osp.join(self.split_root, self.dataset)
            splits = [split]
            for split in splits:
                imgset_file = '{0}_{1}.pth'.format(self.dataset, split)
                imgset_path = osp.join(dataset_path, imgset_file)
                self.images += torch.load(imgset_path)
        
    def exists_dataset(self):
        return osp.exists(osp.join(self.split_root, self.dataset))
    
    def _get_hardpos_verb(self, seg_id, sent_idx):
        """
        Handle the logic for selecting hard positive verb phrases during metric learning.
        Returns the sentence, raw_verb, and tokenized verb if applicable.
        """
        # If the object appears multiple times, no hard positive is used
        if seg_id in self.multi_obj_ref_ids:
            verb_embed = torch.zeros(self.emb_size, dtype=torch.float32)
            return '', verb_embed

        # 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_sent_id = sent_id_list[sent_idx]
            cur_hardpos = hardpos_dict.get(cur_sent_id, {}).get('phrases', [])

        if cur_hardpos:
            # Assign a hard positive verb phrase if available
            rand_index = random.randint(0, len(cur_hardpos) - 1)
            raw_verb = cur_hardpos[rand_index]
            verb_embed = torch.from_numpy(self._get_hardpos_embed(seg_id, cur_sent_id, rand_index))
            # print("Positive phrase : " , raw_verb)
            return raw_verb, verb_embed
        
        verb_embed = torch.zeros(self.emb_size, dtype=torch.float32)
        return '', verb_embed


    def _get_hardpos_embed(self, seg_id, sent_id, rand_index):
        emb_folder = os.path.join(self.all_hp_root, str(seg_id))
        emb_files = sorted([f for f in os.listdir(emb_folder) if f.startswith(f"hp_{sent_id}_") and f.endswith(".npy")])
        selected_emb_file = os.path.join(emb_folder, emb_files[rand_index])

        return np.load(selected_emb_file)


    def pull_item(self, idx):
        # if metric learning and in train mode
        if self.metric_learning and self.augment :
            # sent_idx refers to index of sent among sent_num-1 
            img_file, seg_id, bbox, phrase, sent_idx, sent_num = self.images[idx]
        else :
            img_file, seg_id, bbox, phrase = self.images[idx]
        bbox = np.array(bbox, dtype=int) # x1y1x2y2

        img_path = osp.join(self.im_dir, img_file)
        img = cv2.imread(img_path) # BGR [512, 640, 3]
        ## duplicate channel if gray image
        if img.shape[-1] > 1:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #RGB
        else:
            img = np.stack([img] * 3)
        
        ## seg map
        seg_map = np.load(osp.join(self.mask_root, str(seg_id)+'.npy')) # [512, 640]
        seg_map = np.array(seg_map).astype(np.float32)
        
        if self.metric_learning and self.split == 'train' : 
            return img, phrase, bbox, seg_map, seg_id, sent_idx
        else :
            return img, phrase, bbox, seg_map, seg_id

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        if self.metric_learning and self.augment :
            img, phrase, bbox, seg_map, seg_id, sent_idx = self.pull_item(idx)
        else :
            img, phrase, bbox, seg_map, seg_id = self.pull_item(idx)
            
        phrase = phrase.lower()
        if self.augment:
            augment_flip, augment_hsv, augment_affine, augment_crop, augment_copy, augment_erase = \
                    True,        True,           True,        False,        False,          False

        ## seems a bug in torch transformation resize, so separate in advance
        h,w = img.shape[0], img.shape[1]
        # print("img.shape", img.shape)
        if self.augment:
            ## random horizontal flip
            if augment_flip and random.random() > 0.5:
                img = cv2.flip(img, 1) 
                seg_map = cv2.flip(seg_map, 1) 
                bbox[0], bbox[2] = w-bbox[2]-1, w-bbox[0]-1
                phrase = phrase.replace('right','*&^special^&*').replace('left','right').replace('*&^special^&*','left')

            ## random copy and add left or right
            if augment_copy:
                img, seg_map, phrase, bbox = random_copy(img, seg_map, phrase, bbox)

            ## random erase for occluded
            if augment_erase:
                img, seg_map = random_erase(img, seg_map)

            ## random padding and crop
            if augment_crop:
                img, seg_map = random_crop(img, seg_map, 40, h, w)

            ## random intensity, saturation change
            if augment_hsv:
                fraction = 0.50
                img_hsv = cv2.cvtColor(cv2.cvtColor(img, cv2.COLOR_RGB2BGR), cv2.COLOR_BGR2HSV)
                S = img_hsv[:, :, 1].astype(np.float32)
                V = img_hsv[:, :, 2].astype(np.float32)
                a = (random.random() * 2 - 1) * fraction + 1
                if a > 1:
                    np.clip(S, a_min=0, a_max=255, out=S)
                a = (random.random() * 2 - 1) * fraction + 1
                V *= a
                if a > 1:
                    np.clip(V, a_min=0, a_max=255, out=V)

                img_hsv[:, :, 1] = S.astype(np.uint8)
                img_hsv[:, :, 2] = V.astype(np.uint8)
                img = cv2.cvtColor(cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR), cv2.COLOR_BGR2RGB)

            img, seg_map, ratio, dw, dh = letterbox(img, seg_map, self.imsize)
            bbox[0], bbox[2] = bbox[0]*ratio+dw, bbox[2]*ratio+dw
            bbox[1], bbox[3] = bbox[1]*ratio+dh, bbox[3]*ratio+dh

            ## random affine transformation
            if augment_affine:
                img, seg_map, bbox, M = random_affine(img, seg_map, bbox, \
                    degrees=(-5, 5), translate=(0.10, 0.10), scale=(0.90, 1.10)) # 255 white fill

        else:   ## should be inference, or specified training
            img, _, ratio, dw, dh = letterbox(img, None, self.imsize)
            bbox[0], bbox[2] = bbox[0]*ratio+dw, bbox[2]*ratio+dw
            bbox[1], bbox[3] = bbox[1]*ratio+dh, bbox[3]*ratio+dh

        draw_img = copy.deepcopy(img)
        # Norm, to tensor
        if self.transform is not None:
            img = self.transform(img)
            
                    
        ## encode phrase to clip input
        word_id = clip.tokenize(phrase, 17, truncate=True)
        word_mask = ~ (word_id == 0)
        
        orig_word_id = np.array(word_id, dtype=int)
        orig_word_mask = np.array(word_mask, dtype=int)
                
        # Get hardpos verb phrase
        if self.metric_learning and self.augment:
            raw_hardpos, hardpos_emb = self._get_hardpos_verb(seg_id, sent_idx)
            pos_type = 'nopos'
            if raw_hardpos:
                pos_type = 'hardpos'
                hardpos_id = clip.tokenize(raw_hardpos, self.word_len, truncate=True)
            else:
                # Empty phrase → Create a zero tensor matching shape of tokenized input
                hardpos_id = np.zeros((1, self.word_len), dtype=int)
            
            # Masking
            hardpos_mask = hardpos_id != 0  # Mask should be boolean
            
            hp_word_id = np.array(hardpos_id, dtype=int)
            hp_word_mask = np.array(hardpos_mask, dtype=int)

        if self.augment: # train
            seg_map = cv2.resize(seg_map, (self.imsize // 2, self.imsize // 2),interpolation=cv2.INTER_NEAREST) # (208, 208)
            seg_map = np.reshape(seg_map, [1, np.shape(seg_map)[0], np.shape(seg_map)[1]])
            if self.metric_learning :
                params = {
                    'hp_word_id' : hp_word_id,
                    'hp_word_mask' : hp_word_mask,
                    'hardpos_emb' : hardpos_emb.unsqueeze(0),
                    'pos_type' : pos_type
                }
                return img, orig_word_id, orig_word_mask, np.array(bbox, dtype=np.float32), \
                np.array(seg_map, dtype=np.float32), params
            else : 
                return img, orig_word_id, orig_word_mask, \
                np.array(bbox, dtype=np.float32), np.array(seg_map, dtype=np.float32)
        else:
            seg_map = np.reshape(seg_map, [1, np.shape(seg_map)[0], np.shape(seg_map)[1]])
            return img, orig_word_id, orig_word_mask, \
            np.array(bbox, dtype=np.float32), np.array(seg_map, dtype=np.float32), np.array(ratio, dtype=np.float32), \
            np.array(dw, dtype=np.float32), np.array(dh, dtype=np.float32), self.images[idx][0], self.images[idx][3], np.array(draw_img, dtype=np.uint8)