import copy
import random
import os
import numpy as np
import torch
import json
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
from pycocotools import mask
from pycocotools.coco import COCO
from copy import deepcopy
from typing import Tuple
import cv2
import torch.nn.functional as F
from FCLM.model.language_model.segment_anything.utils.transforms import ResizeLongestSide

def select_ann(coco,img_id, area_min=None, area_max=None):
    cat_ids = coco.getCatIds()
    ann_ids = coco.getAnnIds(imgIds=[img_id], catIds=cat_ids, iscrowd=None)

    if area_min is not None:
        ann_ids = [ann_id for ann_id in ann_ids if coco.anns[ann_id]['area'] >= area_min]

    if area_max is not None:
        ann_ids = [ann_id for ann_id in ann_ids if coco.anns[ann_id]['area'] <= area_max]
    
    return ann_ids

def resize_longest_side(image, target_size):
    """
    调整图像大小，保持纵横比，将最长边缩放到 target_size
    
    参数:
        image: PIL.Image 对象
        target_size: 最长边的目标尺寸（整数）
    
    返回:
        调整后的 PIL.Image 对象
    """
    # 获取原始图像尺寸
    width, height = image.size
    
    # 计算缩放比例
    max_dim = max(width, height)
    scale = target_size / max_dim
    
    # 计算新尺寸（保持整数像素）
    new_width = int(round(width * scale))
    new_height = int(round(height * scale))
    
    # 调整图像大小（使用高质量的 Lanczos 滤镜进行缩放）
    resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    return resized_image
def resize_mask_longest_side(mask_tensor, target_size):
    """
    调整0-1 mask张量的大小，保持纵横比，使用最近邻插值（纯PyTorch实现）
    """
    # 处理维度：确保输入是 (B, C, H, W) 格式（interpolate需要批量维度）
    if len(mask_tensor.shape) == 2:  # (H, W) → (1, 1, H, W)
        mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
    elif len(mask_tensor.shape) == 3:  # (C, H, W) → (1, C, H, W)
        mask_tensor = mask_tensor.unsqueeze(0)
    
    _, _, height, width = mask_tensor.shape
    max_dim = max(height, width)
    scale = target_size / max_dim
    new_height, new_width = int(round(height * scale)), int(round(width * scale))
    
    # 使用 interpolate 进行缩放，mode指定为'nearest'
    resized_mask = F.interpolate(
        mask_tensor,
        size=(new_height, new_width),
        mode='nearest-exact'  # 精确最近邻，避免中间值
    )
    
    # 恢复原始维度（去除批量维度）
    if resized_mask.shape[0] == 1:
        resized_mask = resized_mask.squeeze(0)
    
    return resized_mask


from FCLM.datasets.utils import ANSWER_LIST, SHORT_QUESTION_LIST ,LONG_QUESTION_LIST
QUESTIONS =  [
    'Give me a short description of <region>.',
    'Can you give me a short description of <region>?',
    'Can you provide me with a short description of the region in the picture marked by <region>?',
    "I'm curious about the region represented by <region> in the picture. Could you describe it in few words?",
    'What can you tell me about the region indicated by <region> in the image in few words?',
    "I'd like to know more about the area in the photo labeled <region>. Can you give me a concise description?",
    'Could you describe the region shown as <region> in the picture concisely?',
    'What can you give me about the region outlined by <region> in the photo?',
    'Please provide me with a brief description of the region marked with <region> in the image.',
    'Can you give me a brief introduction of the region labeled as <region> in the picture?',
    "I'm interested in knowing the region represented by <region> in the photo. Can you describe it in several words?",
    'What is the region outlined by <region> in the picture like? Could you give me a streamlined description?',
    'Can you provide me with a brief description of the region in the picture marked by <region>, please?',
    "I'm curious about the region represented by <region> in the picture. Could you describe it in few words, please?",
    'What can you tell me about the region indicated by <region> in the image?',
    "I'd like to know more about the area in the photo labeled <region>, please. Can you give me a simple description?",
    'Could you describe the region shown as <region> in the picture in several words?',
    'Please provide me with a simple description of the region marked with <region> in the image, please.',
    "I'm interested in learning more about the region represented by <region> in the photo. Can you describe it in few words, please?",
    'What is the region outlined by <region> in the picture like, please? Could you give me a simple and clear description?',
    'Please describe the region <region> in the image concisely.',
    'Can you offer a simple analysis of the region <region> in the image?',
    'Could tell me something about the region highlighted by <region> in the picture briefly?',
    'Can you share a simple rundown of the region denoted by <region> in the presented image?'
]

DETAILED_QUESTIONS =  [
    'Can you provide me with a detailed description of the region in the picture marked by <region>?',
    "I'm curious about the region represented by <region> in the picture. Could you describe it in detail?",
    'What can you tell me about the region indicated by <region> in the image?',
    "I'd like to know more about the area in the photo labeled <region>. Can you give me a detailed description?",
    'Could you describe the region shown as <region> in the picture in great detail?',
    'What details can you give me about the region outlined by <region> in the photo?',
    'Please provide me with a comprehensive description of the region marked with <region> in the image.',
    'Can you give me a detailed account of the region labeled as <region> in the picture?',
    "I'm interested in learning more about the region represented by <region> in the photo. Can you describe it in detail?",
    'What is the region outlined by <region> in the picture like? Could you give me a detailed description?',
    'Can you provide me with a detailed description of the region in the picture marked by <region>, please?',
    "I'm curious about the region represented by <region> in the picture. Could you describe it in detail, please?",
    'What can you tell me about the region indicated by <region> in the image, exactly?',
    "I'd like to know more about the area in the photo labeled <region>, please. Can you give me a detailed description?",
    'Could you describe the region shown as <region> in the picture in great detail, please?',
    'What details can you give me about the region outlined by <region> in the photo, please?',
    'Please provide me with a comprehensive description of the region marked with <region> in the image, please.',
    'Can you give me a detailed account of the region labeled as <region> in the picture, please?',
    "I'm interested in learning more about the region represented by <region> in the photo. Can you describe it in detail, please?",
    'What is the region outlined by <region> in the picture like, please? Could you give me a detailed description?',
    'Please describe the region <region> in the image in detail.',
    'Can you offer a thorough analysis of the region <region> in the image?',
    'Could you elaborate on the region highlighted by <region> in the picture provided?',
    'Please share more information about the zone emphasized with <region> in the photo.',
    'What insights can you give ablout the area denoted by <region> in the image presented?',
    'Can you share a comprehensive rundown of the region denoted by <region> in the presented image?',
    "I'd like to know more about the region highlighted by <region> in the picture provided.",
    'Work through the important details of the area <region> in the image.',
    'Illustrate the area represtented by <region> through a descriptive explanation.',
    'Examine the region <region> closely and share its details.'
]

DAM_QUESTIONS = [
    'Give a detailed description of the masked region <region>.',
    'Provide a thorough description of the masked region <region>.'
    'Can you explain the details of the masked area <region>?'
    'Give a detailed account of the masked region <region>.'
    'Describe the masked area <region> comprehensively.'
    'Provide an in-depth description of the masked region <region>.'
    'Explain the specifics of the masked area <region>.'
    'Can you provide a thorough explanation of the masked region <region>?'
    'What are the details of the masked area <region>?'
    'Provide a comprehensive description of the masked area <region>.'
    'What specific details can you provide about the masked region <region>?'
    'Can you give an in-depth account of the masked section <region>?'
    'What are the main characteristics of the masked region <region>?'
    'Give a thorough description of the masked area <region>.'
    'Provide detailed information about the masked area <region>.'
]
class DAM_DATA(Dataset):
    def __init__(self,
                 tokenizer=None,
                 data_args=None,
                 data_path="./data/DAM_data",
                #  max_gt_per_img=5,
                max_gt_per_img=1,
                 ):
        self.PIXEL=1024
        self.data_args = data_args
        self.tokenizer = tokenizer
        self.ann_file_ls = dict()
        self.img_prefix_ls = []
        self.data_type=['COCOStuff','LVIS','Mapillary','OpenImages','PACO','SAM','supp']
        # self.data_type=['COCOStuff','Mapillary','OpenImages','SAM']
        # self.data_type=['COCOStuff','LVIS','Mapillary','OpenImages','PACO'] #org
        # self.data_type=['LVIS','supp'] #debug
        for each in self.data_type:  
            cur_anns_json=os.path.join(data_path,each,'annotations.json')
            cur_dict=json.load(open(cur_anns_json)) 
            prefixed_dict = {f"{each}_{key}": value for key, value in cur_dict.items()}
            # 检查当前字典与已有字典的重复键
            duplicate_keys = self.ann_file_ls.keys() & prefixed_dict.keys()  # 取交集
            if duplicate_keys:
                print(f"警告：数据集 {each} 与已有数据存在重复键：{duplicate_keys}")
            self.ann_file_ls |= prefixed_dict  # 等价于 self.ann_file_ls = self.ann_file_ls | cur_dict
            print(f'DAM DATA:{each} len:{len(prefixed_dict.keys())}')
            cur_img=[]
            if each=='supp':
                for i in range(10):
                    for img in prefixed_dict.keys():
                        img=img.split('_',1)[1]
                        cur_img.append(os.path.join(data_path,each,'images',img)+'.jpg')
            for img in prefixed_dict.keys():
                img=img.split('_',1)[1]
                cur_img.append(os.path.join(data_path,each,'images',img)+'.jpg')
            self.img_prefix_ls.extend(cur_img)

        # assert len(self.img_prefix_ls)==len(self.ann_file_ls.keys())
        self.max_gt_per_img = max_gt_per_img
        self.begin_str = """<image>\nThis provides an overview of the picture.\n"""
        self.transform = ResizeLongestSide(self.PIXEL)
        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
        super().__init__()
    def __len__(self):
        return len(self.img_prefix_ls)
    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 __getitem__(self, idx):
        data_item = self.get_data_item(idx)
        data_dict = self.process_text(data_item=data_item)
        # return data_dict
        # DLC
        # DL-RES
        data_dict_seg = self.process_text_seg(data_item=data_item)
        return data_dict,data_dict_seg
        
    def get_data_item(self, idx):
        
        data_dict={}
        img_path=self.img_prefix_ls[idx]
        image = Image.open(img_path).convert('RGB')
        image=resize_longest_side(image,self.PIXEL)


        ann_key=img_path.split('/')[-3]+'_'+img_path.split('/')[-1].split('.')[0]
        annos=self.ann_file_ls[ann_key]
        mask_ls=[]
        for anno in annos:
            m=mask.decode(anno['mask_rle'])
            mask_ls.append(m)
        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 

        shuffle_ids = torch.randperm(len(mask_ls))
        if len(shuffle_ids) > self.max_gt_per_img:
            shuffle_ids = shuffle_ids[:self.max_gt_per_img]
            
        data_dict['gt_masks'] = mask_ls
        data_dict['category'] =  [anno['category'] for anno in annos]
        data_dict['caption'] = [anno['caption'] for anno in annos]
        data_dict['gt_masks'] = [mask_ls[i] for i in shuffle_ids]
        data_dict['category'] = [data_dict['category'][i] for i in shuffle_ids]
        data_dict['caption'] = [data_dict['caption'][i] for i in shuffle_ids]

        data_dict['img_path']=img_path
        assert len(data_dict['gt_masks']) == len(data_dict['caption']) 
        return data_dict
    def process_text_seg(self, data_item):
        image = data_item['img']
        img_path = data_item['img_path']
        ori_labels = data_item['caption']
        ori_masks = np.array(data_item['gt_masks'])  
        ori_masks = torch.from_numpy(ori_masks) 
        sources = dict()
        sources['conversations'] = []
        for i in range(len(ori_labels)):
            # origin question
            question_template = random.choice(LONG_QUESTION_LIST)
            question=question_template.format(sent=ori_labels[i].lower())
            answer = random.choice(ANSWER_LIST)
            sources['conversations'].append(
                {'from': 'human', 'value': question})
            sources['conversations'].append({'from': 'gpt', 'value': answer})
        
        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')
        # 获取图片当前尺寸
        qwen_img_process = resize_longest_side(qwen_img_process,self.PIXEL)
        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
        #############################
        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])
        data_dict_bs = dict()
        data_dict_bs['input_ids'] = data_dict['input_ids']
        data_dict_bs['labels'] = data_dict['labels']
        masks = ori_masks 
        masks = resize_mask_longest_side(masks,1024)
        sample_num=1
        data_dict_bs['image'] = sam_image.unsqueeze(0).repeat(sample_num,1,1,1)
        data_dict_bs['masks']=masks
        data_dict_bs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
        data_dict_bs['image_grid_thw'] = inputs['image_grid_thw']
        data_dict_bs['img_metas'] = dict()
        data_dict_bs['img_metas']['image'] = torch.tensor(resize).unsqueeze(0).repeat(sample_num,1)
        data_dict_bs['img_metas']['sample_num'] = sample_num
        data_dict_bs['img_metas']['SEG_num'] = sample_num
        return data_dict_bs
    def process_text(self, data_item):
        image = data_item['img']
        img_path = data_item['img_path']
        ori_labels = data_item['caption']
        ori_masks = np.array(data_item['gt_masks'])
        ori_masks = torch.from_numpy(ori_masks) 
        # shuffle_ids = torch.randperm(len(ori_labels))
        # if len(shuffle_ids) > self.max_gt_per_img:
        #     shuffle_ids = shuffle_ids[:self.max_gt_per_img]
        # ori_masks = ori_masks[shuffle_ids]
        # ori_labels = [ori_labels[i] for i in shuffle_ids]
        sources = dict()
        sources['conversations'] = []
        for i in range(len(ori_labels)):
            # origin question
            question = random.choice(DETAILED_QUESTIONS).strip()
            question = question.replace('<region>', '<mask><pos>')
            if i == 0:
                question = self.begin_str + question
            answer = ori_labels[i]
            sources['conversations'].append(
                {'from': 'human', 'value': question})
            sources['conversations'].append({'from': 'gpt', 'value': answer})
        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')
        # 获取图片当前尺寸
        qwen_img_process = resize_longest_side(qwen_img_process,self.PIXEL)

        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
        #############################
        data_dict['image'] = image
        ori_masks = F.interpolate(
            ori_masks.unsqueeze(1),
            size=(image.shape[1], image.shape[2]),
            mode='nearest-exact'  # 精确最近邻，避免中间值
        ).squeeze(1)
        data_dict['masks'] = ori_masks
        data_dict['pixel_values'] = inputs['pixel_values']
        data_dict['image_grid_thw'] = inputs['image_grid_thw']
        return data_dict
class PACO_LVIS_Object_Class_dataset(Dataset):
    def __init__(self,
                tokenizer=None,
                data_args=None,
                data_path="./data/DAM_data",
                max_gt_per_img=5,
                ):

        self.data_args = data_args
        self.tokenizer = tokenizer
        self.ann_file_ls = dict()
        self.img_prefix_ls = []
        self.data_type=['PACO','LVIS']
        for each in self.data_type:
            # cur_anns_json=os.path.join(data_path,each,'annotations.json')
            cur_anns_json=os.path.join(data_path,each,'annotations_new.json')
            cur_dict=json.load(open(cur_anns_json)) 
            prefixed_dict = {f"{each}_{key}": value for key, value in cur_dict.items()}
            # 检查当前字典与已有字典的重复键
            duplicate_keys = self.ann_file_ls.keys() & prefixed_dict.keys()  # 取交集
            
            if duplicate_keys:
                print(f"警告：数据集 {each} 与已有数据存在重复键：{duplicate_keys}")
            self.ann_file_ls |= prefixed_dict  
            print(f'\033[1;34mDAM DATA:{each} len:{len(prefixed_dict.keys())}\033[0m')
            cur_img=[]
            for img in prefixed_dict.keys():
                img=img.split('_',1)[1]
                cur_img.append(os.path.join(data_path,each,'images',img)+'.jpg')
            self.img_prefix_ls.extend(cur_img)
        assert len(self.img_prefix_ls)==len(self.ann_file_ls.keys())
        self.max_gt_per_img = max_gt_per_img
        self.begin_str = """<image>\nThis provides an overview of the picture.\n"""
        super().__init__()
    def __len__(self):
        return len(self.img_prefix_ls)
    def __getitem__(self, idx):
        data_item = self.get_data_item(idx)
        # return data_item
        data_dict = self.process_text(data_item=data_item)
        return data_dict
    def get_data_item(self, idx):
        
        data_dict={}
        img_path=self.img_prefix_ls[idx]
        image = Image.open(img_path).convert('RGB')
        # 获取图片当前尺寸
        width, height = image.size
        
        # max_size = 1024
        # image = resize_longest_side(image, max_size)
        ann_key=img_path.split('/')[-3]+'_'+img_path.split('/')[-1].split('.')[0]
        annos=self.ann_file_ls[ann_key]
        mask_ls=[]
        for anno in annos:
            m=mask.decode(anno['mask_rle'])
            mask_ls.append(m)
        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 
        data_dict['gt_masks'] = mask_ls
        data_dict['category'] =  [anno['category'] for anno in annos]
        data_dict['caption'] = [anno['caption'] for anno in annos]
        data_dict['img_path']=img_path
        assert len(data_dict['gt_masks']) == len(data_dict['caption']) 
        return data_dict
    def process_text(self, data_item):
        image = data_item['img']
        img_path = data_item['img_path']
        
        # ori_labels = data_item['caption']
        ori_masks = np.array(data_item['gt_masks'])
        ori_masks = torch.from_numpy(ori_masks) 
        ori_class = data_item['category']

        shuffle_ids = torch.randperm(len(ori_class))
        if len(shuffle_ids) > self.max_gt_per_img:
            shuffle_ids = shuffle_ids[:self.max_gt_per_img]
        ori_masks = ori_masks[shuffle_ids]
        # ori_labels = [ori_labels[i] for i in shuffle_ids]
        ori_class = [ori_class[i] for i in shuffle_ids]
        sources = dict()

        sources['conversations'] = []
        
        for i in range(len(ori_class)):
            question = 'Please identify the labels of each marked <region> in the image.'
            question = question.replace('<region>', '<mask><pos>')
            if i == 0:
                question = self.begin_str + question
            answer = ori_class[i]
            sources['conversations'].append(
                {'from': 'human', 'value': question})
            sources['conversations'].append({'from': 'gpt', 'value': answer})

        # print(image.shape)
        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')
        # 获取图片当前尺寸
        width, height = qwen_img_process.size
        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
        ############################
        data_dict['image'] = image
        data_dict['masks'] = ori_masks
        data_dict['pixel_values'] = inputs['pixel_values']
        data_dict['image_grid_thw'] = inputs['image_grid_thw']
        return data_dict
if __name__ == '__main__':
    from torch.utils.data import DataLoader
    from tqdm import tqdm
    import transformers
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        './Qwen2.5-VL-3B-Instruct',
        padding_side="right",
        use_fast=True,
    )
    mask_tokens = ['<mask>', '<pos>','<p>','</p>']
    seg_tokens = ['[SEG]']
    tokenizer.add_tokens(mask_tokens, special_tokens=True)
    tokenizer.add_tokens(seg_tokens, special_tokens=True)
    from qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor
    processor = Qwen2_5_VLProcessor.from_pretrained('Qwen2.5-VL-3B-Instruct', trust_remote_code=True)

    dataset = DAM_DATA(tokenizer=tokenizer,data_args=processor)
    total_samples = len(dataset)
    dataloader = DataLoader(
        dataset,
        batch_size=1,  # 如果你不需要批处理，可以保持为 1
        shuffle=True,  # 如果你只是遍历，通常不需要打乱
        num_workers=64,  # 多线程数，根据机器配置调整
        pin_memory=False,  # 如果用到 GPU 可以设为 True 加速
    )
    print(f"开始遍历数据集，共 {total_samples} 个样本...")
    # for i, data in enumerate(tqdm(dataloader, total=total_samples)):
    #     pass  # 你可以在这里对 data 做进一步处理
    for i in range(total_samples):
        try:
            sample = dataset[i]
        except Exception as e:
            break