import os import random import cv2 import numpy as np import torch import torch.nn.functional as F from pycocotools import mask from transformers import CLIPImageProcessor from model.llava import conversation as conversation_lib from model.segment_anything.utils.transforms import ResizeLongestSide, ResizeShortestSide from .grefer import G_REFER from .refer import REFER from .utils import ANSWER_LIST, SHORT_QUESTION_LIST, SINGLE_ANSWER_LIST, MULTI_ANSWER_LIST, EXPAND_QUESTION_LIST class ReferSegDataset(torch.utils.data.Dataset): pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1) pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1) img_size = 1024 ignore_label = 255 def __init__( self, base_image_dir, tokenizer, vision_tower, samples_per_epoch=500 * 8 * 2 * 10, precision: str = "fp32", image_size: int = 224, num_classes_per_sample: int = 3, exclude_val=False, refer_seg_data="refclef||refcoco||refcoco+||refcocog", num_classes_per_question=1, seg_token_num=1, pad_train_clip_images=False, masks_process_with_clip=False, preprocessor_config='', use_expand_question_list=False, ): self.pad_train_clip_images = pad_train_clip_images self.exclude_val = exclude_val self.samples_per_epoch = samples_per_epoch self.num_classes_per_sample = num_classes_per_sample self.base_image_dir = base_image_dir self.image_size = image_size self.tokenizer = tokenizer self.precision = precision self.transform = ResizeLongestSide(image_size) self.short_question_list = SHORT_QUESTION_LIST self.answer_list = ANSWER_LIST self.single_answer_list = SINGLE_ANSWER_LIST self.multi_answer_list = MULTI_ANSWER_LIST self.seg_token_num = seg_token_num self.num_classes_per_question = num_classes_per_question self.masks_process_with_clip = masks_process_with_clip self.pad_train_clip_images = pad_train_clip_images self.clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower) if preprocessor_config == '' else CLIPImageProcessor.from_pretrained(preprocessor_config) self.transform_clip = ResizeLongestSide(self.clip_image_processor.size['shortest_edge']) if use_expand_question_list: self.short_question_list.extend(EXPAND_QUESTION_LIST) if refer_seg_data == 'rs_reason' or refer_seg_data == 'rrsisd': DATA_DIR = base_image_dir else: DATA_DIR = os.path.join(base_image_dir, "refer_seg") self.refer_seg_ds_list = refer_seg_data.split( "||" ) self.refer_seg_data = {} for ds in self.refer_seg_ds_list: if ds == "refcocog": splitBy = "umd" else: splitBy = "unc" if ds == "grefcoco": refer_api = G_REFER(DATA_DIR, ds, splitBy) else: refer_api = REFER(DATA_DIR, ds, splitBy) ref_ids_train = refer_api.getRefIds(split="train") images_ids_train = refer_api.getImgIds(ref_ids=ref_ids_train) refs_train = refer_api.loadRefs(ref_ids=ref_ids_train) refer_seg_ds = {} refer_seg_ds["images"] = [] loaded_images = refer_api.loadImgs(image_ids=images_ids_train) for item in loaded_images: item = item.copy() if ds == "refclef": item["file_name"] = os.path.join( DATA_DIR, "images/saiapr_tc-12", item["file_name"] ) elif ds == 'rs_reason' or ds == 'rrsisd': item["file_name"] = os.path.join( DATA_DIR, "rrsisd/images/rrsisd/JPEGImages", item["file_name"] ) else: item["file_name"] = os.path.join( DATA_DIR, "images/mscoco/images/train2014", item["file_name"] ) refer_seg_ds["images"].append(item) refer_seg_ds["annotations"] = refer_api.Anns print( "dataset {} (refs {}) (train split) has {} images and {} annotations.".format( ds, splitBy, len(refer_seg_ds["images"]), len(refer_seg_ds["annotations"]), ) ) img2refs = {} for ref in refs_train: image_id = ref["image_id"] img2refs[image_id] = img2refs.get(image_id, []) + [ ref, ] refer_seg_ds["img2refs"] = img2refs self.refer_seg_data[ds] = refer_seg_ds def __len__(self): return self.samples_per_epoch def preprocess(self, x: torch.Tensor, decoder_image_size) -> torch.Tensor: """Normalize pixel values and pad to a square input.""" x = (x - self.pixel_mean) / self.pixel_std h, w = x.shape[-2:] padh = decoder_image_size - h padw = decoder_image_size - w x = F.pad(x, (0, padw, 0, padh)) return x def __getitem__(self, idx): ds = random.randint(0, len(self.refer_seg_ds_list) - 1) ds = self.refer_seg_ds_list[ds] refer_seg_ds = self.refer_seg_data[ds] images = refer_seg_ds["images"] annotations = refer_seg_ds["annotations"] img2refs = refer_seg_ds["img2refs"] idx = random.randint(0, len(images) - 1) image_info = images[idx] image_path = image_info["file_name"] image_id = image_info["id"] refs = img2refs[image_id] if len(refs) == 0: return self.__getitem__(0) sents = [] ann_ids = [] for ref in refs: for sent in ref["sentences"]: text = sent["sent"] sents.append(text) ann_ids.append(ref["ann_id"]) max_num_classes_per_sample = self.num_classes_per_question * self.num_classes_per_sample if len(sents) >= max_num_classes_per_sample: sampled_inds = np.random.choice( list(range(len(sents))), size=max_num_classes_per_sample, replace=False ) else: sampled_inds = list(range(len(sents))) sampled_sents = np.vectorize(sents.__getitem__)(sampled_inds).tolist() sampled_ann_ids = [ann_ids[ind] for ind in sampled_inds] sampled_classes = sampled_sents sampled_ann_ids, sampled_classes = allocate_class(sampled_ann_ids, sampled_classes, max_question_num=self.num_classes_per_sample, max_class_per_question=self.num_classes_per_question) image = cv2.imread(image_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) if self.pad_train_clip_images: image_clip = self.transform_clip.apply_image(image) clip_resize = image_clip.shape[:2] image_clip = self.preprocess(torch.from_numpy(image_clip).permute(2, 0, 1).contiguous(), self.clip_image_processor.size['shortest_edge']) else: image_clip = self.clip_image_processor.preprocess(image, return_tensors="pt")[ "pixel_values" ][0] clip_resize = image_clip.shape[-2:] image = self.transform.apply_image(image) resize = image.shape[:2] questions = [] answers = [] seg_token = ["[SEG{}]".format(i) for i in range(self.seg_token_num)] seg_token = ' '.join(seg_token) for text_per_question in sampled_classes: target = '' _seg = [] for i, text in enumerate(text_per_question): assert len(text.split("||")) == 1 if i == len(text_per_question) - 1: _seg.append('[SEG]') if self.seg_token_num == 1 else _seg.append(seg_token) target = target + (' and ' + text) if i != 0 else target + text elif i == 0: target += text _seg.append('[SEG]') if self.seg_token_num == 1 else _seg.append(seg_token) continue else: _seg.append('[SEG]') if self.seg_token_num == 1 else _seg.append(seg_token) target += (', ' + text) if len(_seg) > 1: part1 = ', '.join(_seg[:-1]) part2 = ' and ' + _seg[-1] _seg = part1 + part2 else: _seg = _seg[0] question_template = random.choice(self.short_question_list) questions.append(question_template.format(class_name=target.lower())) separate_answer = random.randint(0, 1) if len(text_per_question) == 1: choice_list = self.answer_list answer_temp = random.choice(choice_list) if self.seg_token_num == 1 else random.choice(choice_list).replace('[SEG]', seg_token) answer_temp = answer_temp.format(class_name=target.lower()) if "{class_name}" in answer_temp else answer_temp answers.append(answer_temp) elif separate_answer: target_answer = [] answer_temp = random.choice(self.single_answer_list) if self.seg_token_num == 1 else random.choice(self.single_answer_list).replace('[SEG]', seg_token) for i, sampled_cls in enumerate(text_per_question): _answer_temp = answer_temp.format(class_name=sampled_cls) if "{class_name}" in answer_temp else answer_temp target_answer.append(_answer_temp[:-1]) if len(target_answer) > 1: part1 = ', '.join(target_answer[:-1]) part2 = ' and ' + target_answer[-1] target_answer = part1 + part2 + '.' else: target_answer = target_answer[0] + '.' answers.append(target_answer) else: answer_temp = random.choice(self.multi_answer_list) _answer_temp = answer_temp.format(class_name=target.lower(), seg=_seg) if "{class_name}" in answer_temp else answer_temp.format(seg=_seg) answers.append(_answer_temp) conversations = [] conv = conversation_lib.default_conversation.copy() i = 0 while i < len(questions): conv.messages = [] conv.append_message(conv.roles[0], questions[i]) conv.append_message(conv.roles[1], answers[i]) conversations.append(conv.get_prompt()) i += 1 image = self.preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous(), self.img_size) flag = False masks = [] for ann_id_per_question in sampled_ann_ids: for ann_id in ann_id_per_question: if isinstance(ann_id, list): flag = True if -1 in ann_id: assert len(ann_id) == 1 m = np.zeros((image_info["height"], image_info["width"])).astype( np.uint8 ) else: m_final = np.zeros( (image_info["height"], image_info["width"]) ).astype(np.uint8) for ann_id_i in ann_id: ann = annotations[ann_id_i] if len(ann["segmentation"]) == 0: m = np.zeros( (image_info["height"], image_info["width"]) ).astype(np.uint8) else: if type(ann["segmentation"][0]) == list: rle = mask.frPyObjects( ann["segmentation"], image_info["height"], image_info["width"], ) else: rle = ann["segmentation"] for i in range(len(rle)): if not isinstance(rle[i]["counts"], bytes): rle[i]["counts"] = rle[i]["counts"].encode() m = mask.decode(rle) m = np.sum( m, axis=2 ) m = m.astype(np.uint8) m_final = m_final | m m = m_final masks.append(m) continue ann = annotations[ann_id] if len(ann["segmentation"]) == 0: m = np.zeros((image_info["height"], image_info["width"])).astype( np.uint8 ) masks.append(m) continue if type(ann["segmentation"][0]) == list: rle = mask.frPyObjects( ann["segmentation"], image_info["height"], image_info["width"] ) else: rle = ann["segmentation"] for i in range(len(rle)): if not isinstance(rle[i]["counts"], bytes): rle[i]["counts"] = rle[i]["counts"].encode() m = mask.decode(rle) m = np.sum( m, axis=2 ) m = m.astype(np.uint8) masks.append(m) masks = np.stack(masks, axis=0) masks = torch.from_numpy(masks) label = torch.ones(masks.shape[1], masks.shape[2]) * self.ignore_label if self.masks_process_with_clip: mask_shape = image_clip.shape[-1] if len(masks) == 0: masks = torch.zeros(0, mask_shape, mask_shape) else: masks = transform_mask(masks, mask_shape) return ( image_path, image, image_clip, conversations, masks, label, resize, clip_resize, questions, sampled_classes, ) def allocate_class(sampled_ann_ids, sampled_ann_classes, max_question_num=3, max_class_per_question=3): if len(sampled_ann_ids) < max_question_num: max_question_num = len(sampled_ann_ids) sample_num = len(sampled_ann_classes) question_id = np.arange(max_question_num) class_counts = np.arange(max_question_num) * 0 new_sampled_ann_ids = [[] for _ in range(max_question_num)] new_sampled_ann_classes = [[] for _ in range(max_question_num)] sample_ids = np.arange(sample_num) np.random.shuffle(sample_ids) for i in range(sample_num): if 0 in class_counts: choose_id = np.random.choice(np.where(class_counts == 0)[0], size=1)[0] else: choose_id = np.random.choice(np.where(class_counts < max_class_per_question)[0], size=1)[0] class_counts[choose_id] += 1 sample_id = sample_ids[i] new_sampled_ann_ids[choose_id].append(sampled_ann_ids[sample_id]) new_sampled_ann_classes[choose_id].append(sampled_ann_classes[sample_id]) return new_sampled_ann_ids, new_sampled_ann_classes def transform_mask(masks, size): height, width = masks.shape[-2:] short, long = (width, height) if width <= height else (height, width) requested_new_short = size new_short, new_long = requested_new_short, int(requested_new_short * long / short) new_shape = (new_long, new_short) if width <= height else (new_short, new_long) masks = F.interpolate(masks[None].float(), size=new_shape, mode="nearest")[0].bool() orig_height, orig_width = new_shape crop_height, crop_width = size, size crop_height, crop_width = int(crop_height), int(crop_width) top = (orig_height - crop_height) // 2 bottom = top + crop_height left = (orig_width - crop_width) // 2 right = left + crop_width assert top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width masks = masks[..., top:bottom, left:right] return masks def center_crop_image(image, size): orig_height, orig_width = image.shape[:2] crop_height, crop_width = size, size crop_height, crop_width = int(crop_height), int(crop_width) top = (orig_height - crop_height) // 2 bottom = top + crop_height left = (orig_width - crop_width) // 2 right = left + crop_width assert top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width image = image[top:bottom, left:right] return image