import torch import os from enum import Enum from tqdm import tqdm import numpy as np from detectron2.structures import BitMasks from psalm.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, \ DEFAULT_IM_END_TOKEN, DEFAULT_SEG_TOKEN, SEG_TOKEN_INDEX from psalm.model.builder import load_pretrained_model from psalm.utils import disable_torch_init from psalm.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria import cv2 from torch.utils.data import Dataset, DataLoader from psalm import conversation as conversation_lib from psalm.train.train_datasets_eval import COCO_interactive_dataset_extrametric # debug from detectron2.structures import BoxMode from detectron2.data import MetadataCatalog, DatasetCatalog from typing import Dict, Optional, Sequence, List from dataclasses import dataclass, field import torch.distributed as dist import transformers from pathlib import Path from psalm.eval.segmentation_evaluation import openseg_classes from natsort import natsorted COLOR_MAP = openseg_classes.ADE20K_150_CATEGORIES import re from psalm.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, SEG_TOKEN_INDEX, CLS_TOKEN_INDEX, REGION_TOKEN_INDEX, REFER_TOKEN_INDEX import json from pycocotools import mask as mask_utils # 定义命令行参数 @dataclass class DataArguments: data_path: str = field(default=None, metadata={"help": "Path to the training data."}) lazy_preprocess: bool = False is_multimodal: bool = False image_folder: Optional[str] = field(default='/path/to/val2017') model_path: Optional[str] = field(default="/path/to/model") mask_config: Optional[str] = field(default="./psalm/mask_config/maskformer2_swin_base_384_bs16_50ep.yaml") image_aspect_ratio: str = 'square' image_grid_pinpoints: Optional[str] = field(default=None) json_path: str = '/path/to/coco' model_map_name: str = 'psalm_video' version: str = 'llava_phi' segmentation: bool = True eval_batch_size: int = 1 # debug dataloader_num_workers: int = 8 seg_task: Optional[str] = field(default="region") region_mask_type: Optional[str] = field(default=None) with_memory: bool = False resume: bool = False using_autocast: bool = False resume_path: Optional[str] = field(default=None) save_format: Optional[str] = field(default=None) save_path: Optional[str] = field(default=None) #定义collect函数 @dataclass class DataCollatorForCOCODatasetV2(object): """Collate examples for supervised fine-tuning.""" tokenizer: transformers.PreTrainedTokenizer #sequence表示列表、元组等有序对象，instances的类型表示为字典组成的有序列表，其中一个字典表示一帧图像 def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]: input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels")) input_ids = torch.nn.utils.rnn.pad_sequence( input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id) labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX) input_ids = input_ids[:, :self.tokenizer.model_max_length] labels = labels[:, :self.tokenizer.model_max_length] batch = dict( input_ids=input_ids, labels=labels, attention_mask=input_ids.ne(self.tokenizer.pad_token_id), ) if 'image' in instances[0]: images = [instance['image'] for instance in instances] if all(x is not None and x.shape == images[0].shape for x in images): batch['images'] = torch.stack(images) else: batch['images'] = images if 'vp_image' in instances[0]: vp_images = [instance['vp_image'] for instance in instances] if all(x is not None and x.shape == vp_images[0].shape for x in vp_images): batch['vp_images'] = torch.stack(vp_images) else: batch['vp_images'] = vp_images for instance in instances: for key in ['input_ids', 'labels', 'image']: del instance[key] batch['seg_info'] = [instance for instance in instances] if 'dataset_type' in instances[0]: batch['dataset_type'] = [instance['dataset_type'] for instance in instances] if 'class_name_ids' in instances[0]: class_name_ids = [instance['class_name_ids'] for instance in instances] if any(x.shape != class_name_ids[0].shape for x in class_name_ids): batch['class_name_ids'] = torch.nn.utils.rnn.pad_sequence( class_name_ids, batch_first=True, padding_value=-1, ) else: batch['class_name_ids'] = torch.stack(class_name_ids, dim=0) if 'token_refer_id' in instances[0]: token_refer_id = [instance['token_refer_id'] for instance in instances] batch['token_refer_id'] = token_refer_id if 'cls_indices' in instances[0]: cls_indices = [instance['cls_indices'] for instance in instances] if any(x.shape != cls_indices[0].shape for x in cls_indices): batch['cls_indices'] = torch.nn.utils.rnn.pad_sequence( cls_indices, batch_first=True, padding_value=-1, ) else: batch['cls_indices'] = torch.stack(cls_indices, dim=0) if 'random_idx' in instances[0]: random_idxs = [instance['random_idx'] for instance in instances] batch['random_idx'] = torch.stack(random_idxs, dim=0) if 'class_name_embedding_indices' in instances[0]: class_name_embedding_indices = [instance['class_name_embedding_indices'] for instance in instances] class_name_embedding_indices = torch.nn.utils.rnn.pad_sequence( class_name_embedding_indices, batch_first=True, padding_value=0) batch['class_name_embedding_indices'] = class_name_embedding_indices if 'refer_embedding_indices' in instances[0]: refer_embedding_indices = [instance['refer_embedding_indices'] for instance in instances] refer_embedding_indices = torch.nn.utils.rnn.pad_sequence( refer_embedding_indices, batch_first=True, padding_value=0) batch['refer_embedding_indices'] = refer_embedding_indices return batch # 定义处理model输出结果的函数 def parse_outputs(outputs,gt_mask): res_list = [] for output in outputs: # gt = output['gt'].cpu().numpy().astype(np.uint8) pred_mask = output['instances'].pred_masks pred_mask = pred_mask.cpu().numpy() scores = output['instances'].scores.transpose(1,0).cpu().numpy() gt_mask = output['gt'].cpu().numpy().astype(np.uint8) try: pred_cls = output['instances'].pred_classes.cpu().numpy() except: pred_cls = None assert scores.shape[0] == gt_mask.shape[0] for i in range(gt_mask.shape[0]): res = { 'pred':pred_mask, 'gt': gt_mask[i], 'scores':scores[i], 'pred_cls':pred_cls } res_list.append(res) return res_list # 定义Dataset类 class DAVIS_Dataset(COCO_interactive_dataset_extrametric): #注意，这里所有的处理逻辑针对的都是一帧图像 def __getitem__(self, idx): data = self.data[idx] #图片的相对路径名称，like2017/trainval/JPEGImages/480p/bike-packing/00001.jpg image_file = data['image'] #image_folder是data_root根路径 在这里是data_segswap image_folder = self.data_args.image_folder data_dict = {} #file_name是图片的完整路径名称，like /data/Davis/2017/trainval/JPEGImages/480p/bike-packing/00001.jpg data_dict['file_name'] = os.path.join(image_folder, image_file) data_dict['height'] = data['image_info']['height'] data_dict['width'] = data['image_info']['width'] #image_id可以理解为计数器，编号 data_dict['image_id'] = data['new_img_id'] #annotations，本帧对应的注释，coco格式的分割mask，一张图片可能包含多个实例的mask data_dict['annotations'] = data['anns'] #vp_annotations，每段视频中第一帧的注释 data_dict['vp_annotations'] = data['first_frame_anns'] #vp_image，每段视频中第一帧的完整路径，like /data/Davis/2017/trainval/JPEGImages/480p/bike-packing/00000.jpg data_dict['vp_image'] = os.path.join(image_folder,data['first_frame_image']) for annotation in data_dict['annotations']: annotation['bbox_mode'] = BoxMode.XYXY_ABS #边界框左上角和右下角的坐标都为原点，意思是将边界框置为空框 annotation['bbox'] = [0,0,0,0] annotation['image_id'] = data['new_img_id'] for annotation in data_dict['vp_annotations']: annotation['bbox_mode'] = BoxMode.XYXY_ABS annotation['bbox'] = [0,0,0,0] annotation['image_id'] = data['new_img_id'] processor = self.data_args.image_processor['null_mask'] # debug：处理null mask #尝试从命令行参数中获取region_mask_type region_mask_type = getattr(self.data_args,'region_mask_type',None) if region_mask_type is not None: region_mask_type = region_mask_type.split('||') #print("region_mask_type:", region_mask_type) #根据region_mask_type和mask_format（这里是0、1掩码），对原始的data_dict进行预处理，将Detectron2格式的dataset dict转化为MaskFormer格式的 data_dict = processor.preprocess(data_dict,region_mask_type=region_mask_type,mask_format='bitmask') #num_target，本帧图像中有多少个对象 #下面的一小段代码，主要是利用llama处理输入的文本，生成对应的token num_target = len(data_dict['instances']) # 是一个特殊的占位符，表示图像的输入 prefix_inst = 'This is an image , Please segment by given regions' # 占位符来表示每个需要分割的区域，用逗号分隔，最后一个 以句号结束，例如，如果有 3 个区域，结果是 ' , , .' regions_inst = ' ,' * (num_target - 1) + ' .' sources_value = f'\nThis is all regions: {regions_inst}\n' #sources构建了一个人类和模型交互的对话格式，定义了来自人类的输入和来自模型的输出 sources = [ [{'from': 'human', 'value': prefix_inst + sources_value}, {'from': 'gpt', 'value': '\n[SEG]'}]] text_dict = self.preprocess_llama2(sources, self.tokenizer) #input_ids是模型的实际输入，是由分词器将文本 sources 转换为的一系列数字标识（token IDs） input_ids = text_dict['input_ids'][0] #labels是模型训练时的token的真实标签，与input_ids对应 labels = text_dict['labels'][0] data_dict['input_ids'] = input_ids data_dict['labels'] = labels data_dict['dataset_type'] = 'region_coco' return data_dict def evaluation(): # 模型的加载 parser = transformers.HfArgumentParser(DataArguments) data_args = parser.parse_args_into_dataclasses()[0] disable_torch_init() model_path = os.path.expanduser(data_args.model_path) model_name = get_model_name_from_path(model_path) print(f'current model is {model_path}') tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, model_args=data_args, mask_config=data_args.mask_config, device='cuda') data_args.image_processor = image_processor data_args.is_multimodal = True conversation_lib.default_conversation = conversation_lib.conv_templates[data_args.version] # 数据集的加载 data_collator = DataCollatorForCOCODatasetV2(tokenizer=tokenizer) dataloader_params = { "batch_size": data_args.eval_batch_size, "num_workers": data_args.dataloader_num_workers, } def load_ref_dataset(): return DAVIS_Dataset(json_path=data_args.json_path, tokenizer=tokenizer, data_args=data_args) #注册load_ref_dataset函数，方便快速获取数据集 DatasetCatalog.register('refcoco_dataset', load_ref_dataset) MetadataCatalog.get('refcoco_dataset').set(stuff_classes=['object'],) # 模型导入到device device = 'cuda' if torch.cuda.is_available() else 'cpu' if data_args.using_autocast: model.to(device=device).eval() # debug：不指定模型精度 else: model.to(device=device,dtype=torch.float).eval() # 定义处理takes的范围 gt_json_path = "/home/yuqian_fu/Projects/PSALM/gt_test.json" with open(gt_json_path, "r") as fp: gt_json = json.load(fp) save_path_json = "/scratch/yuqian_fu/report_soccer_new.json" # debug: 实验前修改 data_path = "/home/yuqian_fu/Projects/PSALM/egoexo_test_framelevel.json" # debug: 实验前修改 with open(data_path, "r") as fp: datas = json.load(fp) splits_path = "/home/yuqian_fu/Projects/ego-exo4d-relation/correspondence/split.json" with open(splits_path, "r") as fp: splits = json.load(fp) takes_all = splits["test"] # 随机从takes_all中选取20个，并设置随机种子，可以复现 # np.random.seed(0) # takes_all = np.random.choice(takes_all, 20, replace=False) json_path = "/home/yuqian_fu/Projects/PSALM/filter_takes_byname_test.json" with open(json_path, "r") as fp: datas_tmp = json.load(fp) takes_all = datas_tmp['soccer'] # debug: 实验前修改 # NUM = len(takes_all) // 4 # takes_all = takes_all[:NUM] #takes_all = takes_all[NUM:NUM*2] #takes_all = takes_all[NUM*2:NUM*3] #takes_all = takes_all[NUM*3:NUM*4] #takes_all = takes_all[NUM*4:NUM*5] #takes_all = takes_all[NUM*5:NUM*6] # takes_all_2 = takes_all[NUM*7:NUM*9] #takes_all = takes_all_1 + takes_all_2 #takes_all = takes_all[NUM*6:NUM*7] #takes_all = takes_all[NUM*7:NUM*8] #takes_all = takes_all[NUM*8:NUM*9] #takes_all = takes_all[NUM*9:NUM*10] #takes_all = takes_all[NUM*10:NUM*11] #takes_all = takes_all[NUM*11:NUM*12] #takes_all = takes_all[NUM*12:NUM*13] #takes_all = takes_all[NUM*13:NUM*14] #takes_all = takes_all[NUM*14:NUM*15] #takes_all = takes_all[NUM*15:] # num_tmp = len(takes_all)//2 # takes_all = takes_all[num_tmp:] #takes_all = takes_all[NUM*7:] #takes_all = ["8c952699-0c25-453b-92dd-52b0580248db"] # debug: 测试单个take # resume机制 if data_args.resume: with open(data_args.resume_path, "r") as fp: result = json.load(fp) # 删除result字典中最后处理的take_id # last_processed_take = "xxx" # del result[last_processed_take] processed_takes = set(result.keys()) takes_all = [take_id for take_id in takes_all if take_id not in processed_takes] else: result = {} # 混合精度推理 if data_args.using_autocast: scaler = torch.cuda.amp.autocast(enabled=True) # 记录target注释文件缺失的数量 anno_miss_num = 0 with torch.no_grad(): for take_id in tqdm(takes_all): print("current take_id:", take_id) # 获取针对每个take的标注文件 with open(f'{data_args.image_folder}/{take_id}/annotation.json', 'r') as fp: annotations = json.load(fp) gt = gt_json["annotations"][take_id] # 获取每个take下的所有物体，并创建从fill_number到物体名称的映射 objs = natsorted(list(annotations["masks"].keys())) #debug: 是否有必要使用natsort coco_id_to_cont_id = {cont_id + 1: coco_id for cont_id, coco_id in enumerate(objs)} id_range = list(coco_id_to_cont_id.keys()) # 筛选出所有video-name为take_id的数据，构造数据集 datas_list = [] for data in datas: if data['video_name'] == take_id: datas_list.append(data) # print("len(datas_list):", len(datas_list)) # debug eval_dataset = DAVIS_Dataset(datas_list, tokenizer=tokenizer, data_args=data_args) eval_dataloader = DataLoader(eval_dataset, batch_size=dataloader_params['batch_size'], collate_fn=data_collator, num_workers=dataloader_params['num_workers']) # 保存每个take下的结果 pred_json = {'masks': {}, 'subsample_idx': annotations['subsample_idx']} objs_after = [] # debug: 统计推理之后obj的个数 for idx, inputs in enumerate(eval_dataloader): # 准备inputs inputs = {k: v.to(device) if torch.is_tensor(v) else v for k, v in inputs.items()} h = inputs['seg_info'][0]['height'] w = inputs['seg_info'][0]['width'] # 提取cam target_cam = inputs['seg_info'][0]['file_name'].split('/')[-2] query_cam = inputs['seg_info'][0]['vp_file_path'].split('/')[-2] # debug pair_key = f'{query_cam}_{target_cam}' # print("query_cam:", query_cam) # debug # print("target_cam:", target_cam) # 提取id id = inputs['seg_info'][0]['file_name'].split('/')[-1].split('.')[0] # print("id:", id) # debug # 看能否跳过加速 flag = False for i in range(len(inputs['seg_info'][0]['instances'].vp_fill_number)): cur_fill_number_tmp = inputs['seg_info'][0]['instances'].vp_fill_number[i] obj_name = coco_id_to_cont_id[cur_fill_number_tmp.item()] objs_after.append(obj_name) if int(id) not in gt["annotated_frames"][obj_name][target_cam]: flag = True cur_pred = np.zeros((h, w), dtype=np.uint8) #debug: 改路径 save_path = f'{data_args.save_path}/{take_id}/{target_cam}/{obj_name}/{id}.png' os.makedirs(os.path.dirname(save_path), exist_ok=True) cv2.imwrite(save_path, cur_pred.astype(np.uint8)) # 1) 保证第一层 obj_name 存在 if obj_name not in pred_json['masks']: pred_json['masks'][obj_name] = {} # 2) 保证第二层 pair_key 存在 if pair_key not in pred_json['masks'][obj_name]: pred_json['masks'][obj_name][pair_key] = {} pred_json['masks'][obj_name][f'{query_cam}_{target_cam}'][id] = {'pred_mask': save_path, 'confidence': 1.0} break if flag: #print(f"flag is True, skipping take_id {take_id}, id {id}") continue # 能否加速-v2 flag_tmp = False for i in range(len(inputs['seg_info'][0]['instances'].vp_fill_number)): cur_fill_number_tmp = inputs['seg_info'][0]['instances'].vp_fill_number[i] obj_name = coco_id_to_cont_id[cur_fill_number_tmp.item()] objs_after.append(obj_name) if id not in annotations["masks"][obj_name][target_cam].keys(): flag_tmp = True cur_pred = np.zeros((h, w), dtype=np.uint8) #debug: 改路径 save_path = f'{data_args.save_path}/{take_id}/{target_cam}/{obj_name}/{id}.png' os.makedirs(os.path.dirname(save_path), exist_ok=True) cv2.imwrite(save_path, cur_pred.astype(np.uint8)) # 1) 保证第一层 obj_name 存在 if obj_name not in pred_json['masks']: pred_json['masks'][obj_name] = {} # 2) 保证第二层 pair_key 存在 if pair_key not in pred_json['masks'][obj_name]: pred_json['masks'][obj_name][pair_key] = {} pred_json['masks'][obj_name][f'{query_cam}_{target_cam}'][id] = {'pred_mask': save_path, 'confidence': 0.0} break if flag_tmp: #print(f"flag is True, skipping take_id {take_id}, id {id}") continue # 使用混合精度 if data_args.using_autocast: with torch.cuda.amp.autocast(): outputs = model.eval_video( input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], images=inputs['images'].float(), vp_images=inputs['vp_images'].float(), seg_info=inputs['seg_info'], labels=inputs['labels'] ) # 不使用混合精度 else: outputs = model.eval_video( input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], images=inputs['images'].float(), vp_images=inputs['vp_images'].float(), seg_info=inputs['seg_info'], labels=inputs['labels'] ) if torch.cuda.is_available(): torch.cuda.synchronize() output = outputs[0] pred_mask = output['instances'].pred_masks pred_mask = pred_mask.cpu().numpy() scores = output['instances'].scores.transpose(1, 0).cpu().numpy() gt_mask = output['gt'].cpu().numpy().astype(np.uint8) assert len(scores) == len(inputs['seg_info'][0]['instances'].vp_fill_number) prev_idx = [] for i in range(len(scores)): cur_scores = scores[i] cur_fill_number = inputs['seg_info'][0]['instances'].vp_fill_number[i] # debug：如果填充物体id不在所有物体的索引列表中，跳过 if cur_fill_number not in id_range: print(f"cur_fill_number {cur_fill_number} not in id_range, skipping...") raise ValueError(f"cur_fill_number {cur_fill_number} not in id_range, skipping...") max_score, idx = torch.topk(torch.tensor(cur_scores), 10, largest=True, sorted=True) idx = idx.cpu().numpy() for i in range(10): if idx[i] not in prev_idx: prev_idx.append(idx[i]) pick_idx = idx[i] pick_score = max_score[i] break # 根据cur_fill_number逆映射找到obj-name obj_name = coco_id_to_cont_id[cur_fill_number.item()] objs_after.append(obj_name) # debug: 统计推理之后obj的个数 # 对获取到的obj_name进行合法性筛查 if target_cam not in annotations['masks'][obj_name].keys(): print(f"target_cam {target_cam} not in {obj_name}, skipping...") raise ValueError(f"target_cam {target_cam} not in {obj_name}, skipping...") if id not in annotations["masks"][obj_name][target_cam].keys(): anno_miss_num += 1 # print(f"id {id} not in {target_cam}, skipping...") # 编码cur_pred // 保存mask图片 cur_pred = pred_mask[pick_idx, :] if data_args.save_format == 'rle': cur_pred = mask_utils.encode(np.asfortranarray(cur_pred.astype(np.uint8))) cur_pred['counts'] = cur_pred['counts'].decode('ascii') elif data_args.save_format == 'png': save_path = f'{data_args.save_path}/{take_id}/{target_cam}/{obj_name}/{id}.png' # debug: 改路径 os.makedirs(os.path.dirname(save_path), exist_ok=True) cv2.imwrite(save_path, cur_pred.astype(np.uint8)) else: raise ValueError(f"Unsupported save format: {data_args.save_format}") # 1) 保证第一层 obj_name 存在 if obj_name not in pred_json['masks']: pred_json['masks'][obj_name] = {} # 2) 保证第二层 pair_key 存在 if pair_key not in pred_json['masks'][obj_name]: pred_json['masks'][obj_name][pair_key] = {} if data_args.save_format == 'rle': pred_json['masks'][obj_name][f'{query_cam}_{target_cam}'][id] = {'pred_mask': cur_pred, 'confidence': pick_score.item()} elif data_args.save_format == 'png': pred_json['masks'][obj_name][f'{query_cam}_{target_cam}'][id] = {'pred_mask': save_path, 'confidence': pick_score.item()} #检查一下pred_json['masks']的内容是否为空 if len(pred_json['masks']) == 0: print(f"pred_json['masks'] is empty for take_id {take_id}, skipping...") # continue # bug # 将这个take下的所有结果存储 #存储之前，确保take下每个物体都存在，有的物体下没有targer_cam，需要置空写死 check_obj = set(objs) - set(objs_after) if len(check_obj) > 0: for obj in check_obj: # 对于缺失的物体，分两种情况。一种是obj下没有任何cam，另一种是obj下有cam但是没有ids cams = annotations['masks'][obj].keys() exo_cams = [x for x in cams if 'aria' not in x] ego_cams = [x for x in cams if 'aria' in x] # 如果exo_cams和ego_cams都非空，则需要增加一个cam的键 if len(exo_cams) > 0 and len(ego_cams) > 0: pred_json['masks'][obj] = {} ego = ego_cams[0] for exo in exo_cams: pred_json['masks'][obj][f"{exo}_{ego}"] = {} else: print(f"{take_id}缺失物体{obj}") pred_json['masks'][obj] = {} #pred_json['masks'][obj] = "xxx" result[take_id] = pred_json # 每个take定期保存，防止中断 # with open(save_path_json, 'w') as fp: # json.dump(result, fp) # 保存最后的结果 with open(save_path_json, "w") as fp: json.dump(result, fp) # 打印miss anno样本的数目 print(f"Total number of missing annotations: {anno_miss_num}") if __name__ == '__main__': evaluation() # path = "/scratch/yuqian_fu/competition_test_20250516_single_take_complete.json" # with open(path, "r") as fp: # result = json.load(fp) # pred = result["40dc3bbc-c8c4-4e6d-a2a5-32a357f3c291"] # print(len(list(pred['masks'].keys())))