# import copy
# import random
# import os
# import numpy as np
# import torch
# import ast
# from FCLM.train.train import preprocess, preprocess_multimodal
# from PIL import Image
# from torch.utils.data import Dataset, DataLoader
# import pandas as pd
# from torchvision import transforms
# import glob
# import json
# import cv2
# import torch.nn.functional as F
# DEFAULT_IMAGE_TOKEN = "<image>"
# SHORT_QUESTION_LIST = [
#     DEFAULT_IMAGE_TOKEN + "\n" + "Can you segment the {class_name} in this image?",
#     DEFAULT_IMAGE_TOKEN + "\n" + "Please segment the {class_name} in this image.",
#     DEFAULT_IMAGE_TOKEN
#     + "\n"
#     + "What is {class_name} in this image? Please respond with segmentation mask.",
#     DEFAULT_IMAGE_TOKEN
#     + "\n"
#     + "What is {class_name} in this image? Please output segmentation mask.",
# ]

# LONG_QUESTION_LIST = [
#     DEFAULT_IMAGE_TOKEN + "\n" + "{sent} Please respond with segmentation mask.",
#     DEFAULT_IMAGE_TOKEN + "\n" + "{sent} Please output segmentation mask.",
# ]

# EXPLANATORY_QUESTION_LIST = [
#     "Please output segmentation mask and explain why.",
#     "Please output segmentation mask and explain the reason.",
#     "Please output segmentation mask and give some explanation.",
# ]

# ANSWER_LIST = [
#     "It is [SEG].",
#     "Sure, [SEG].",
#     "Sure, it is [SEG].",
#     "Sure, the segmentation result is [SEG].",
#     "[SEG].",
# ]


# def get_mask_from_json(json_path, img):
#     try:
#         with open(json_path, "r") as r:
#             anno = json.loads(r.read())
#     except:
#         with open(json_path, "r", encoding="cp1252") as r:
#             anno = json.loads(r.read())

#     inform = anno["shapes"]
#     comments = anno["text"]
#     is_sentence = anno["is_sentence"]

#     height, width = img.shape[:2]

#     ### sort polies by area
#     area_list = []
#     valid_poly_list = []
#     for i in inform:
#         label_id = i["label"]
#         points = i["points"]
#         if "flag" == label_id.lower():  ## meaningless deprecated annotations
#             continue

#         tmp_mask = np.zeros((height, width), dtype=np.uint8)
#         cv2.polylines(tmp_mask, np.array([points], dtype=np.int32), True, 1, 1)
#         cv2.fillPoly(tmp_mask, np.array([points], dtype=np.int32), 1)
#         tmp_area = tmp_mask.sum()

#         area_list.append(tmp_area)
#         valid_poly_list.append(i)

#     ### ground-truth mask
#     sort_index = np.argsort(area_list)[::-1].astype(np.int32)
#     sort_index = list(sort_index)
#     sort_inform = []
#     for s_idx in sort_index:
#         sort_inform.append(valid_poly_list[s_idx])

#     mask = np.zeros((height, width), dtype=np.uint8)
#     for i in sort_inform:
#         label_id = i["label"]
#         points = i["points"]

#         if "ignore" in label_id.lower():
#             label_value = 255  # ignored during evaluation
#         else:
#             label_value = 1  # target

#         cv2.polylines(mask, np.array([points], dtype=np.int32), True, label_value, 1)
#         cv2.fillPoly(mask, np.array([points], dtype=np.int32), label_value)

#     return mask, comments, is_sentence
# class REASON_SEG(Dataset):
#     def __init__(self,
#                 tokenizer=None,
#                 data_args=None,
#                 base_image_dir=None,
#                 samples_per_epoch=500 * 8 * 2 * 10,
#                 precision: str = "fp32",
#                 # image_size: int = 224,
#                 num_classes_per_sample: int = 1,
#                 exclude_val=False,
#                 reason_seg_data="ReasonSeg|train",
#                 explanatory=0.1,
#                  ):
#         self.data_args = data_args
#         self.tokenizer = tokenizer
#         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.precision = precision
#         # self.transform = ResizeLongestSide(image_size)
#         # self.clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower)

#         self.short_question_list = SHORT_QUESTION_LIST
#         self.long_question_list = LONG_QUESTION_LIST
#         self.answer_list = ANSWER_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, 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,
#                     "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"],
#                 }
        
#         super().__init__()
#     def __len__(self):
#         return self.samples_per_epoch
#     def __getitem__(self, idx):
#         data_item = self.get_data_item(idx)
#         return data_item
#     def preprocess_pad(self, x,H,W) -> torch.Tensor:
#         """Normalize pixel values and pad to a square input."""
#         # Pad
#         h, w = x.shape[-2:]
#         padh = H - h
#         padw = W - w
#         x = F.pad(x, (0, padw, 0, padh))
#         return x
#     def get_data_item(self, idx):
#         data_dict={}
        
#         images, jsons = self.reason_seg_data
#         idx = random.randint(0, len(images) - 1)
#         img_path = images[idx]
#         json_path = jsons[idx]
#         image = Image.open(img_path).convert('RGB')
        
#         mask, sents, is_sentence = get_mask_from_json(json_path, cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB))
#         normalize = transforms.Compose([
#             transforms.ToTensor(),                      # 转为Tensor并自动归一化到[0,1]
#             transforms.Normalize(
#                 mean=[0.485, 0.456, 0.406],            # ImageNet均值
#                 std=[0.229, 0.224, 0.225]              # ImageNet标准差
#             )
#         ])
#         image = normalize(image) 
#         data_dict['img'] = 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_name = img_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 = []
#         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()))

#             # add explanation if applicable
#             img_name = img_path.split("/")[-1]
#             if self.explanatory != -1 and img_name in self.img_to_explanation:
#                 if choice == 0:  # [SEG] token
#                     answers.append(random.choice(self.answer_list))
#                     #['It is [SEG].', 'Sure, [SEG].', 'Sure, it is [SEG].', 'Sure, the segmentation result is [SEG].', '[SEG].']
#                 elif choice == 1:  # [SEG] token + text answer
#                     image_name = img_path.split("/")[-1]
#                     answer = self.img_to_explanation[image_name]["outputs"]
#                     answer = random.choice(self.answer_list) + " {}".format(answer)
#                     questions[-1] = (
#                         DEFAULT_IMAGE_TOKEN
#                         + "\n"
#                         + text
#                         + " {}".format(random.choice(self.explanatory_question_list))
#                     )
#                     #'<image>\nWhat is the object that is typically used during the hurdling event in track and field? Please output segmentation mask and give some explanation.'
#                     answers.append(answer)
#                     #"It is [SEG]. In the picture, there is a hurdle, which is typically used in the hurdling event of track and field competitions. The athlete in the image is jumping over this hurdle during the competition. Hurdles are designed to challenge the athletes' agility, speed, and jumping ability as they attempt to clear them while running at high speeds."
#                 elif choice == 2:  # vanilla text answer
#                     image_name = img_path.split("/")[-1]
#                     answer = self.img_to_explanation[image_name]["outputs"]
#                     questions[-1] = DEFAULT_IMAGE_TOKEN + "\n" + text
#                     #'<image>\nWhat is the object that is typically used during the hurdling event in track and field?'
#                     answers.append(answer)
#                     #"In the picture, there is a hurdle, which is typically used in the hurdling event of track and field competitions. The athlete in the image is jumping over this hurdle during the competition. Hurdles are designed to challenge the athletes' agility, speed, and jumping ability as they attempt to clear them while running at high speeds."
#                 else:
#                     raise ValueError("Not implemented yet.")
#             else:
#                 answers.append(random.choice(self.answer_list))

#             conversations = []
#             sources = dict()
#             sources['conversations'] = []
#             for i in range(len(questions)):
#                 if i>0:
#                     sources['conversations'].append(
#                         {'from': 'human', 'value': questions[i].replace('<image>\n', '')})
#                 else:
#                     sources['conversations'].append(
#                         {'from': 'human', 'value': questions[i]})
#                 sources['conversations'].append({'from': 'gpt', 'value': answers[i]})

        
#         sources = preprocess_multimodal(
#             copy.deepcopy([sources['conversations']]),
#             self.data_args,
#             0,if_qwen=True)
        
#         text=sources[0][0]['value']
#         processor=self.data_args.image_processor
#         qwen_img_process=Image.open(img_path).convert('RGB')
#         W,H=qwen_img_process.size
#         qwen_img_process=qwen_img_process.resize((W//2,H//2))
#         inputs = processor(
#             text=[text],
#             images=qwen_img_process,
#             videos=None,
#             padding=True,
#             return_tensors="pt",
#         )
#         tmp=inputs['input_ids']
#         N=len(tmp[tmp==self.tokenizer('<|image_pad|>')['input_ids'][0]]) 
        
#         data_dict = preprocess(
#             sources,
#             self.tokenizer,
#             has_image=True
#             )
#         ####repeat <|image_pad|> * N
#         old_data_dict=data_dict
#         insert_position=torch.where(old_data_dict['input_ids']==self.tokenizer('<|image_pad|>')['input_ids'][0])[1]
#         img_padding=torch.tensor([self.tokenizer('<|image_pad|>')['input_ids'][0]]*N).unsqueeze(0)
#         target_padding=torch.tensor([-100]*N).unsqueeze(0)
#         # (1,N) <|image_pad|>
#         new_inputs=torch.cat((old_data_dict['input_ids'][:,:insert_position],img_padding,old_data_dict['input_ids'][:,insert_position+1:]),dim=1)
#         # [1, 435] 345 个 image_pad
#         new_target=torch.cat((old_data_dict['labels'][:,:insert_position],target_padding,old_data_dict['labels'][:,insert_position+1:]),dim=1)
#         data_dict['input_ids']=new_inputs
#         data_dict['labels']=new_target
#         #############################
#         # import ipdb;ipdb.set_trace()
        
#         # get single
#         if isinstance(i, int):
#             data_dict = dict(input_ids=data_dict['input_ids'][0],
#                              labels=data_dict['labels'][0])
#         H,W=image.shape[1],image.shape[2]
#         MAX_H,MAX_W=1024,1024 #SAM grounding
#         # import ipdb;ipdb.set_trace()
#         # image_pad=self.preprocess_pad(image,H,W)
#         image=torch.nn.functional.interpolate(image.unsqueeze(0),
#                                                 # size=(H//2,W//2),
#                                                 size=(H,W),
#                                                 mode='bilinear',
#                                                 align_corners=False).squeeze(0)
#         data_dict['image'] = image
#         masks = np.stack(sampled_masks, axis=0)
#         masks = torch.from_numpy(masks)
#         masks=torch.nn.functional.interpolate(masks.float().unsqueeze(0),
#                                                 # size=(H//2,W//2),
#                                                 size=(H,W),
#                                                 mode='nearest'
#                                                 ).squeeze(0)
#         data_dict['masks'] = masks
#         data_dict['pixel_values'] = inputs['pixel_values']
#         data_dict['image_grid_thw'] = inputs['image_grid_thw']

#         return data_dict


import copy
import random
import os
import numpy as np
import torch
import ast
from FCLM.train.train import preprocess, preprocess_multimodal
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from torchvision import transforms
import glob
import json
import cv2
import threading
import torch.nn.functional as F
from FCLM.model.language_model.segment_anything.utils.transforms import ResizeLongestSide
cv2.setNumThreads(0)
DEFAULT_IMAGE_TOKEN = "<image>"
SHORT_QUESTION_LIST = [
    DEFAULT_IMAGE_TOKEN + "\n" + "Can you segment the {class_name} in this image?",
    DEFAULT_IMAGE_TOKEN + "\n" + "Please segment the {class_name} in this image.",
    DEFAULT_IMAGE_TOKEN
    + "\n"
    + "What is {class_name} in this image? Please respond with segmentation mask.",
    DEFAULT_IMAGE_TOKEN
    + "\n"
    + "What is {class_name} in this image? Please output segmentation mask.",
]

LONG_QUESTION_LIST = [
    DEFAULT_IMAGE_TOKEN + "\n" + "{sent} Please respond with segmentation mask.",
    DEFAULT_IMAGE_TOKEN + "\n" + "{sent} Please output segmentation mask.",
]

EXPLANATORY_QUESTION_LIST = [
    "Please output segmentation mask and explain why.",
    "Please output segmentation mask and explain the reason.",
    "Please output segmentation mask and give some explanation.",
]

ANSWER_LIST = [
    "It is [SEG].",
    "Sure, [SEG].",
    "Sure, it is [SEG].",
    "Sure, the segmentation result is [SEG].",
    "[SEG].",
]


def get_mask_from_json(json_path, img):
    try:
        with open(json_path, "r") as r:
            anno = json.loads(r.read())
    except:
        with open(json_path, "r", encoding="cp1252") as r:
            anno = json.loads(r.read())

    inform = anno["shapes"]
    comments = anno["text"]
    is_sentence = anno["is_sentence"]

    height, width = img.shape[:2]

    ### sort polies by area
    area_list = []
    valid_poly_list = []
    for i in inform:
        label_id = i["label"]
        points = i["points"]
        if "flag" == label_id.lower():  ## meaningless deprecated annotations
            continue

        tmp_mask = np.zeros((height, width), dtype=np.uint8)
        cv2.polylines(tmp_mask, np.array([points], dtype=np.int32), True, 1, 1)
        cv2.fillPoly(tmp_mask, np.array([points], dtype=np.int32), 1)
        tmp_area = tmp_mask.sum()

        area_list.append(tmp_area)
        valid_poly_list.append(i)

    ### ground-truth mask
    sort_index = np.argsort(area_list)[::-1].astype(np.int32)
    sort_index = list(sort_index)
    sort_inform = []
    for s_idx in sort_index:
        sort_inform.append(valid_poly_list[s_idx])

    mask = np.zeros((height, width), dtype=np.uint8)
    for i in sort_inform:
        label_id = i["label"]
        points = i["points"]

        if "ignore" in label_id.lower():
            label_value = 255  # ignored during evaluation
        else:
            label_value = 1  # target

        cv2.polylines(mask, np.array([points], dtype=np.int32), True, label_value, 1)
        cv2.fillPoly(mask, np.array([points], dtype=np.int32), label_value)

    return mask, comments, is_sentence
class REASON_SEG(Dataset):
    
    def __init__(self,
                tokenizer=None,
                data_args=None,
                base_image_dir=None,
                samples_per_epoch=500 * 2 * 10 * 2,
                precision: str = "fp32",
                # image_size: int = 224,
                num_classes_per_sample: int = 1,
                exclude_val=False,
                reason_seg_data="ReasonSeg|train",
                explanatory=-1,
                 ):
        self.pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
        self.pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
        self.img_size = 1024
        self.data_args = data_args
        self.tokenizer = tokenizer
        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.precision = precision
        self.transform = ResizeLongestSide(self.img_size)
        self.short_question_list = SHORT_QUESTION_LIST
        self.long_question_list = LONG_QUESTION_LIST
        self.answer_list = ANSWER_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, 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,
                    "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"],
                }
        super().__init__()
    def __len__(self):
        return self.samples_per_epoch
    def __getitem__(self, idx):
        data_item = self.get_data_item(idx)
        return data_item
    def preprocess_pad(self, x: torch.Tensor,if_img=True) -> torch.Tensor:
        """Normalize pixel values and pad to a square input."""
        # Normalize colors
        if if_img:
            x = (x - self.pixel_mean) / self.pixel_std

        # Pad
        h, w = x.shape[-2:]
        padh = self.img_size - h
        padw = self.img_size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x
    def get_data_item(self, idx):
        data_dict={}

        images, jsons = self.reason_seg_data
        # idx = random.randint(0, len(images) - 1)
        idx = idx % len(images)
        img_path = images[idx]
        json_path = jsons[idx]
        image = Image.open(img_path).convert('RGB')
        
        mask, sents, is_sentence = get_mask_from_json(json_path, cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB))
        normalize = transforms.Compose([
            transforms.ToTensor(),                      # 转为Tensor并自动归一化到[0,1]
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],            # ImageNet均值
                std=[0.229, 0.224, 0.225]              # ImageNet标准差
            )
        ])
        image = normalize(image) 
        data_dict['img'] = 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_name = img_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)
        else:
            choice = 0
        questions = []
        answers = []
        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()))

            # add explanation if applicable
            img_name = img_path.split("/")[-1]
            if self.explanatory != -1 and img_name in self.img_to_explanation:
                if choice == 0:  # [SEG] token
                    answers.append(random.choice(self.answer_list))
                    #['It is [SEG].', 'Sure, [SEG].', 'Sure, it is [SEG].', 'Sure, the segmentation result is [SEG].', '[SEG].']
                elif choice == 1:  # [SEG] token + text answer
                    image_name = img_path.split("/")[-1]
                    answer = self.img_to_explanation[image_name]["outputs"]
                    answer = random.choice(self.answer_list) + " {}".format(answer)
                    questions[-1] = (
                        DEFAULT_IMAGE_TOKEN
                        + "\n"
                        + text
                        + " {}".format(random.choice(self.explanatory_question_list))
                    )
                    #'<image>\nWhat is the object that is typically used during the hurdling event in track and field? Please output segmentation mask and give some explanation.'
                    answers.append(answer)
                    #"It is [SEG]. In the picture, there is a hurdle, which is typically used in the hurdling event of track and field competitions. The athlete in the image is jumping over this hurdle during the competition. Hurdles are designed to challenge the athletes' agility, speed, and jumping ability as they attempt to clear them while running at high speeds."
                elif choice == 2:  # vanilla text answer
                    image_name = img_path.split("/")[-1]
                    answer = self.img_to_explanation[image_name]["outputs"]
                    questions[-1] = DEFAULT_IMAGE_TOKEN + "\n" + text
                    #'<image>\nWhat is the object that is typically used during the hurdling event in track and field?'
                    answers.append(answer)
                    #"In the picture, there is a hurdle, which is typically used in the hurdling event of track and field competitions. The athlete in the image is jumping over this hurdle during the competition. Hurdles are designed to challenge the athletes' agility, speed, and jumping ability as they attempt to clear them while running at high speeds."
                else:
                    raise ValueError("Not implemented yet.")
            else:
                answers.append(random.choice(self.answer_list))

            conversations = []
            sources = dict()
            sources['conversations'] = []
            for i in range(len(questions)):
                if i>0:
                    sources['conversations'].append(
                        {'from': 'human', 'value': questions[i].replace('<image>\n', '')})
                else:
                    sources['conversations'].append(
                        {'from': 'human', 'value': questions[i]})
                sources['conversations'].append({'from': 'gpt', 'value': answers[i]})
            # print('sources:',sources)
            # import ipdb;ipdb.set_trace()
        sources = preprocess_multimodal(
            copy.deepcopy([sources['conversations']]),
            self.data_args,
            0,if_qwen=True)
        
        text=sources[0][0]['value']
        processor=self.data_args.image_processor
        qwen_img_process=Image.open(img_path).convert('RGB')
        W,H=qwen_img_process.size
        qwen_img_process=qwen_img_process.resize((W//4,H//4))
        # qwen_img_process=qwen_img_process
        inputs = processor(
            text=[text],
            images=qwen_img_process,
            videos=None,
            padding=True,
            return_tensors="pt",
        )
        tmp=inputs['input_ids']
        N=len(tmp[tmp==self.tokenizer('<|image_pad|>')['input_ids'][0]]) 
        
        data_dict = preprocess(
            sources,
            self.tokenizer,
            has_image=True
            )
        ####repeat <|image_pad|> * N
        old_data_dict=data_dict
        insert_position=torch.where(old_data_dict['input_ids']==self.tokenizer('<|image_pad|>')['input_ids'][0])[1]
        img_padding=torch.tensor([self.tokenizer('<|image_pad|>')['input_ids'][0]]*N).unsqueeze(0)
        target_padding=torch.tensor([-100]*N).unsqueeze(0)
        # (1,N) <|image_pad|>
        new_inputs=torch.cat((old_data_dict['input_ids'][:,:insert_position],img_padding,old_data_dict['input_ids'][:,insert_position+1:]),dim=1)
        # [1, 435] 345 个 image_pad
        new_target=torch.cat((old_data_dict['labels'][:,:insert_position],target_padding,old_data_dict['labels'][:,insert_position+1:]),dim=1)
        data_dict['input_ids']=new_inputs #torch.Size([1, 1628])
        data_dict['labels']=new_target #torch.Size([1, 1628])
        #############################
        # 
        
        # get single
        if isinstance(i, int):
            data_dict = dict(input_ids=data_dict['input_ids'][0],
                             labels=data_dict['labels'][0])
        sam_image = self.transform.apply_image(cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB))
        resize = sam_image.shape[:2]
        sam_image = self.preprocess_pad(torch.from_numpy(sam_image).permute(2, 0, 1).contiguous())
        data_dict['image'] = sam_image #torch.Size([3, 1024, 1024])
        # sam_mask = self.transform.apply_image(np.stack(sampled_masks, axis=0).transpose(1,2,0))
        # sam_mask = self.preprocess_pad(torch.from_numpy(sam_mask).unsqueeze(0).contiguous(),if_img=False)
        # sam_mask[sam_mask>0]=1
        masks = np.stack(sampled_masks, axis=0)
        masks = torch.from_numpy(masks)
        masks=torch.nn.functional.interpolate(masks.unsqueeze(0),
                                                size=(masks.shape[1]//4,masks.shape[2]//4),
                                                mode='nearest'
                                                ).squeeze(0)
        data_dict['masks'] = masks #torch.Size([1, 966, 1288])
        data_dict['pixel_values'] = inputs['pixel_values'] #torch.Size([6256, 1176])
        data_dict['image_grid_thw'] = inputs['image_grid_thw'] #tensor([[ 1, 68, 92]])
        data_dict['img_metas'] = resize #(768, 1024)
        # import ipdb;ipdb.set_trace()
        return data_dict
