import glob import json import os import random 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 import copy from model.segment_anything.utils.transforms import ResizeLongestSide from model.llava import conversation as conversation_lib from .utils import ( 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, ) 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 MultiReasonSegValDataset(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, val_dataset, image_size=1024, seg_token_num=1, pad_val_clip_images=False, masks_process_with_clip=False, preprocessor_config='', crop_sam_image=False ): self.pad_val_clip_images= pad_val_clip_images self.masks_process_with_clip = masks_process_with_clip self.base_image_dir = base_image_dir self.image_size = image_size self.tokenizer = tokenizer self.transform = ResizeLongestSide(image_size) 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']) self.short_question_list = SHORT_QUESTION_LIST self.long_question_list = LONG_QUESTION_LIST self.answer_list = ANSWER_LIST reason_seg_data, split = val_dataset.split("|") assert split == 'val' print(base_image_dir) json_file_name = "./dataset/muse_val.json" with open(json_file_name, 'r') as f: reason_file = json.load(f) images = [] anns = [] questions = [] answers = [] self.reason_seg_data = reason_file print("number of reason_seg samples: ", len(images)) def __len__(self): return len(self.reason_seg_data) 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): 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]) segs = image_info['ann_list'] question = image_info['questions'] gt_answer = image_info['answers'] gt_target_count = [] gt_category_name = [] name_list = [ann['rephrased_name'] if 'rephrased_name' in ann else ann['category_name'] for ann in segs] _name_list = [] name_count = {} for name in name_list: if name not in name_count: name_count[name] = 1 else: name_count[name] += 1 max_name_count = copy.deepcopy(name_count) name_loc = [] phrase_loc = [] for name, ann in zip(name_list, segs): x, y, w, h = ann['bbox'] x0 = x x1 = x + w y0 = y y1 = y + h bbox_str = str([x0, y0, x1, y1]) if max_name_count[name] == 1: _name_list.append(name) name_loc.append('{} at {}'.format(name, bbox_str)) else: name_loc.append('{} {} at {}'.format(name, str(max_name_count[name] - name_count[name] + 1), bbox_str)) _name_list.append('{} {}'.format(name, str(max_name_count[name] - name_count[name] + 1))) name_count[name] -= 1 name_loc = ', '.join(name_loc) name_str = ', '.join(_name_list) prompt_ins = "These objects in the image and their respective bounding box coordinates are as follows: {}. The image height is {}, width is {}.".format(name_loc, image_info['height'], image_info['width']) img = cv2.imread(image_path) images = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) ori_size = images.shape[:2] if self.pad_val_clip_images: images_clip = self.transform_clip.apply_image(images) clip_resize = images_clip.shape[:2] images_clip = self.preprocess(torch.from_numpy(images_clip).permute(2, 0, 1).contiguous(), self.clip_image_processor.size['shortest_edge']) else: images_clip = self.clip_image_processor.preprocess(images, return_tensors="pt")[ "pixel_values" ][0] clip_resize = images_clip.shape[:2] images = self.transform.apply_image(images) resize = images.shape[:2] masks = [] if len(segs) == 0: return self[0] for answer_list in gt_answer: gt_target_count.append(len(answer_list)) gt_category_name.append(['(' + ann['rephrased_name'] + ' ' + str([ann['bbox'][0], ann['bbox'][1], ann['bbox'][0]+ann['bbox'][2], ann['bbox'][1]+ann['bbox'][3]]) + ')' for ann in answer_list]) 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) sampled_sents = question sampled_answers = gt_answer sampled_masks = masks image_name = image_path.split("/")[-1] questions = [] answers = [] for text, answer in zip(sampled_sents, sampled_answers): question_template = random.choice(self.long_question_list) questions.append(question_template.format(sent=text)) 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], "") 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 = images_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, images_clip, conversations, masks, label, resize, clip_resize, (questions, gt_target_count, gt_category_name, prompt_ins), sampled_sents, False, True ) 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 = MultiReasonSegValDataset("data", tokenizer, "openai/clip-vit-large-patch14", val_dataset="MultiReasonseg|val") for i in range(len(dataset)): data = dataset[i]