import glob import json import os import random import cv2 import numpy as np import torch import torch.nn.functional as F from transformers import CLIPImageProcessor from model.llava import conversation as conversation_lib from model.segment_anything.utils.transforms import ResizeLongestSide, ResizeShortestSide from .data_processing import get_mask_from_json from .utils import (ANSWER_LIST, DEFAULT_IMAGE_TOKEN, EXPLANATORY_QUESTION_LIST, LONG_QUESTION_LIST, SHORT_QUESTION_LIST, SINGLE_ANSWER_LIST, MULTI_ANSWER_LIST, EXPAND_QUESTION_LIST, EXPAND_LONG_QUESTION_LIST) class ReasonSegDataset(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="ReasonSeg|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 = 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) self.long_question_list.extend(EXPAND_LONG_QUESTION_LIST) reason_seg_data, splits = reason_seg_data.split("|") splits = splits.split("_") images = [] for split in splits: images_split = glob.glob( os.path.join( base_image_dir, "reason_seg", reason_seg_data, split, "*.jpg" ) ) images.extend(images_split) jsons = [path.replace(".jpg", ".json") for path in images] self.reason_seg_data = (images, jsons) print("number of reason_seg samples: ", len(images)) if explanatory != -1: self.explanatory_question_list = EXPLANATORY_QUESTION_LIST self.img_to_explanation = {} with open( os.path.join( base_image_dir, "reason_seg", reason_seg_data, "explanatory", "train.json", ) ) as f: items = json.load(f) for item in items: img_name = item["image"] self.img_to_explanation[img_name] = { "query": item["query"], "outputs": item["outputs"], } print("len(self.img_to_explanation): ", len(self.img_to_explanation)) 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): images, jsons = self.reason_seg_data idx = random.randint(0, len(images) - 1) image_path = images[idx] json_path = jsons[idx] image = cv2.imread(image_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) ori_size = image.shape[:2] 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:] mask, sents, is_sentence = get_mask_from_json(json_path, image) if len(sents) >= self.num_classes_per_sample: sampled_inds = np.random.choice( list(range(len(sents))), size=self.num_classes_per_sample, replace=False ) else: sampled_inds = list(range(len(sents))) sampled_sents = np.vectorize(sents.__getitem__)(sampled_inds).tolist() sampled_masks = [ (mask == 1).astype(np.float32) for _ in range(len(sampled_inds)) ] image = self.transform.apply_image(image) resize = image.shape[:2] image_name = image_path.split("/")[-1] if self.explanatory != -1 and image_name in self.img_to_explanation: if random.random() < self.explanatory: choice = 2 else: choice = random.randint(0, 1) questions = [] answers = [] seg_token = ["[SEG{}]".format(i) for i in range(self.seg_token_num)] seg_token = ' '.join(seg_token) for text in sampled_sents: if is_sentence: question_template = random.choice(self.long_question_list) questions.append(question_template.format(sent=text)) else: question_template = random.choice(self.short_question_list) questions.append(question_template.format(class_name=text.lower())) img_name = image_path.split("/")[-1] if self.explanatory != -1 and img_name in self.img_to_explanation: if choice == 0: answer_temp = random.choice(self.answer_list) if self.seg_token_num == 1 else random.choice(self.answer_list).replace('[SEG]', seg_token) answers.append(answer_temp) elif choice == 1: image_name = image_path.split("/")[-1] answer = self.img_to_explanation[image_name]["outputs"] answer_temp = random.choice(self.answer_list) if self.seg_token_num == 1 else random.choice(self.answer_list).replace('[SEG]', seg_token) answer = answer_temp + " {}".format(answer) questions[-1] = ( DEFAULT_IMAGE_TOKEN + "\n" + text + " {}".format(random.choice(self.explanatory_question_list)) ) answers.append(answer) elif choice == 2: image_name = image_path.split("/")[-1] answer = self.img_to_explanation[image_name]["outputs"] questions[-1] = DEFAULT_IMAGE_TOKEN + "\n" + text answers.append(answer) else: raise ValueError("Not implemented yet.") else: answer_temp = random.choice(self.answer_list) if self.seg_token_num == 1 else random.choice(self.answer_list).replace('[SEG]', seg_token) answers.append(answer_temp) 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 image = self.preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous(), self.img_size) image_name = image_path.split("/")[-1] if ( self.explanatory != -1 and image_name in self.img_to_explanation and choice == 2 ): masks = torch.rand(0, *ori_size) label = torch.ones(ori_size) * self.ignore_label else: 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, image, image_clip, conversations, masks, label, resize, clip_resize, questions, sampled_sents, ) 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