# -*- 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)