import glob import json import os import random from unicodedata import category from pycocotools import mask import cv2 import numpy as np import torch import torch.nn.functional as F from transformers import CLIPImageProcessor, PretrainedConfig import transformers from model.segment_anything.utils.transforms import ResizeLongestSide from model.llava import conversation as conversation_lib from .utils import ( MR_SINGLE_ANSWER_LIST, MR_MULTI_ANSWER_LIST, ANSWER_LIST, DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IMAGE_TOKEN, EXPLANATORY_QUESTION_LIST, LONG_QUESTION_LIST, SHORT_QUESTION_LIST, EXPAND_LONG_QUESTION_LIST, ) from transformers.image_utils import make_list_of_images, to_numpy_array, infer_channel_dimension_format from transformers.image_transforms import convert_to_rgb, to_channel_dimension_format from transformers.image_processing_utils import get_size_dict class MultiReasonSegDataset(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, reason_seg_data="MultiReasonSeg|train", explanatory=0.1, 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.exclude_val = exclude_val self.reason_seg_data = reason_seg_data self.samples_per_epoch = samples_per_epoch self.explanatory = explanatory 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.long_question_list = LONG_QUESTION_LIST self.answer_list = ANSWER_LIST self.single_answer_list = MR_SINGLE_ANSWER_LIST self.multi_answer_list = MR_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.long_question_list.extend(EXPAND_LONG_QUESTION_LIST) print("___________self.single_answer_list:", self.single_answer_list) print("___________self.multi_answer_list:", self.multi_answer_list) reason_seg_data, split = reason_seg_data.split("|") json_file_name = './dataset/muse_train.json' with open(json_file_name, 'r') as f: reason_file = json.load(f) images = [] anns = [] questions = [] answers = [] self.reason_seg_data = reason_file 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): idx = random.randint(0, len(self.reason_seg_data) - 1) image_info = self.reason_seg_data[idx] if 'file_name' in image_info: image_root = os.path.join(self.base_image_dir, "refer_seg/images/mscoco/images/train2014") image_path = os.path.join(image_root, image_info['file_name']) else: if 'train2017' in image_info['coco_url']: image_root = os.path.join(self.base_image_dir, "refer_seg/images/mscoco/images/train2017") image_path = os.path.join(image_root, image_info['coco_url'].split('/')[-1]) else: image_root = os.path.join(self.base_image_dir, "refer_seg/images/mscoco/images/val2017") image_path = os.path.join(image_root, image_info['coco_url'].split('/')[-1]) anns = image_info['ann_list'] question = image_info['questions'] if 'questions' in image_info else None gt_answer = image_info['answers'] if 'answers' in image_info else None if question is not None: text_answers = image_info['text_answers'] if 'text_answers' in image_info else [None] * len(gt_answer) else: text_answers = None img = cv2.imread(image_path) images = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) ori_size = images.shape[:2] if self.pad_train_clip_images: image_clip = self.transform_clip.apply_image(images) 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(images, return_tensors="pt")[ "pixel_values" ][0] clip_resize = image_clip.shape[-2:] images = self.transform.apply_image(images) resize = images.shape[:2] masks = [] if len(anns) == 0: return self[0] category_ids = [ann['category_id'] for ann in anns] category_ids = list(set(category_ids)) sampled_num = min(self.num_classes_per_sample, len(category_ids)) sampled_category_ids = np.random.choice(category_ids, size=sampled_num, replace=False) sampled_sents = question sampled_answers = gt_answer sampled_masks = masks sample_text_answers = text_answers image_name = image_path.split("/")[-1] questions = [] answers = [] use_assign_list = [] seg_token = ["[SEG{}]".format(i) for i in range(self.seg_token_num)] seg_token = ' '.join(seg_token) if question is not None: for text, answer_list, text_answer in zip(sampled_sents, sampled_answers, sample_text_answers): question_template = random.choice(self.long_question_list) questions.append(question_template.format(sent=text)) for answer in answer_list: rle = mask.frPyObjects(answer["segmentation"], image_info["height"], image_info["width"]) m = mask.decode(rle) if len(m.shape) > 2: m = np.sum(m, axis=2) m = m.astype(np.uint8) masks.append(m) use_assign = False if text_answer is not None: if text_answer.count('{seg}') != len(answer_list): return self[0] _text_answer = text_answer.format(seg='[SEG]') if self.seg_token_num == 1 else text_answer.format(seg=seg_token) answers.append(_text_answer) use_assign_list.append(False) else: target_list = [a['rephrased_name'] if (random.random() > 0.1 and 'rephrased_name' in a) else a['category_name'] for a in answer_list ] target_answer = [] separate_answer = random.randint(0, 1) _seg = ['[SEG]'] * len(target_list) if len(target_list) > 1: part1 = ', '.join(_seg[:-1]) part2 = ' and ' + _seg[-1] _seg = part1 + part2 else: _seg = _seg[0] if separate_answer: choice_list = self.single_answer_list answer_temp = random.choice(choice_list) if self.seg_token_num == 1 else random.choice(choice_list).replace('[SEG]', seg_token) use_assign = False if "{class_name}" in answer_temp else True for i, sampled_cls in enumerate(target_list): _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] + '.' else: answer_temp = random.choice(self.multi_answer_list) _answer_temp = answer_temp.format(class_name=', '.join(target_list).lower(), seg=_seg) if "{class_name}" in answer_temp else answer_temp.format(seg=_seg) use_assign = False if "{class_name}" in answer_temp else True _answer_temp = _answer_temp if self.seg_token_num == 1 else _answer_temp.replace('[SEG]', seg_token) target_answer = _answer_temp answers.append(target_answer) use_assign_list.append(use_assign) else: for sampled_category_id in sampled_category_ids: question_template = random.choice(self.instance_question_list) category_names = self.lvis_name_dict[str(sampled_category_id)] category_name = random.choice(category_names) questions.append(question_template.format(class_name=category_name)) answer_list = [ann for ann in anns if ann['category_id'] == sampled_category_id] for answer in answer_list: rle = mask.frPyObjects(answer["segmentation"], image_info["height"], image_info["width"]) m = mask.decode(rle) if len(m.shape) > 2: m = np.sum(m, axis=2) m = m.astype(np.uint8) masks.append(m) target_list = [a['rephrased_name'] if random.random() > 0.1 else a['category_name'] for a in answer_list ] target_answer = [] separate_answer = random.randint(0, 1) _seg = ['[SEG]'] * len(target_list) if len(target_list) > 1: part1 = ', '.join(_seg[:-1]) part2 = ' and ' + _seg[-1] _seg = part1 + part2 else: _seg = _seg[0] separate_answer = random.randint(0, 1) choice_list = self.single_answer_list answer_temp = random.choice(choice_list) if self.seg_token_num == 1 else random.choice(choice_list).replace('[SEG]', seg_token) use_assign = False if "{class_name}" in answer_temp else True for i, sampled_cls in enumerate(target_list): _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) use_assign_list.append(use_assign) conversations = [] conv = conversation_lib.default_conversation.copy() roles = {"human": conv.roles[0], "gpt": conv.roles[1]} 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 images = self.preprocess(torch.from_numpy(images).permute(2, 0, 1).contiguous(), self.img_size) image_name = image_path.split("/")[-1] masks = np.stack(sampled_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, images, image_clip, conversations, masks, label, resize, clip_resize, questions, sampled_sents, use_assign_list ) 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 clip_image_process(clip_image_processor, images): images = make_list_of_images(images) images = [convert_to_rgb(image) for image in images] images = [to_numpy_array(image) for image in images] input_data_format = infer_channel_dimension_format(images[0]) resize_transform = ResizeLongestSide(clip_image_processor.size['shortest_edge']) images = [resize_transform.apply_image(image) for image in images] images = [ clip_image_processor.rescale(image=image, scale=clip_image_processor.rescale_factor) for image in images] images = [ clip_image_processor.normalize(image=image, mean=clip_image_processor.image_mean, std=clip_image_processor.image_std) for image in images ] images = [ F.pad(torch.tensor(image).permute(2, 0, 1), (0, 224-image.shape[1], 0, 224-image.shape[0])) for image in images ] return images[0] if __name__ == "__main__": version = "checkpoints/llava-v1.6-vicuna-7b" tokenizer = transformers.AutoTokenizer.from_pretrained( version, cache_dir=None, model_max_length=512, padding_side="right", use_fast=False, ) tokenizer.pad_token = tokenizer.unk_token num_added_tokens = tokenizer.add_tokens("[SEG]") ret_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids dataset = MultiReasonSegDataset("data", tokenizer, "openai/clip-vit-large-patch14") for i in range(len(dataset)): import pdb;pdb.set_trace() data = dataset[i]