import json import os from PIL import Image import numpy as np from pycocotools.mask import encode, decode, frPyObjects from tqdm import tqdm import copy from natsort import natsorted import cv2 #这一脚本的主要目的是把ego mask缩放到exo size，ego的图片大小还是704 704 if __name__ == '__main__': root_path = '/data/work2-gcp-europe-west4-a/yuqian_fu/Ego/data_segswap' #跑实验需改动 save_path = os.path.join(root_path, 'egoexo_val_exosize.json') #获取takes_id split_path = "/home/yuqian_fu/Projects/ego-exo4d-relation/correspondence/SegSwap/data/split.json" with open(split_path, "r") as fp: data_split = json.load(fp) val_set = data_split["val"] #跑实验需改动 # val_set = ["1d0f3c10-ed0a-4f60-b0d2-a516690ff1cf"] #用来计数 new_img_id = 0 #用来存储json中的数据 egoexo_dataset = [] ''' build_DAVIS.py的代码逻辑是先处理每个视频的第一帧，第一帧中的unique_instances、高宽等信息用于该视频下后续的每一帧。 注意，unique_instances代表的是第一帧下像素的所有类别信息，如果该视频下后续的帧中有像素的类别不在unique_instances中，会报错 ''' bad_case = [] for val_name in tqdm(val_set): #不同视角下两个相机的总路径 vid_root_path = os.path.join(root_path, val_name) anno_path = os.path.join(vid_root_path, "annotation.json") with open(anno_path, 'r') as fp: annotations = json.load(fp) #取出本take下的所有物体 # objs = list(annotations["masks"].keys()) # 确保每次obj的顺序一样，构造的coco_id_to_cont_id数字与物体的映射一样 objs = natsorted(list(annotations["masks"].keys())) print("the total obj num are:", len(objs)) print(f"objs:{objs}") #将物体名称映射为id "cook":1 从1开始，区别于背景 #TODO看看这个要不要修改为以obj_ref中的物体为准 coco_id_to_cont_id = {coco_id: cont_id+1 for cont_id, coco_id in enumerate(objs)} #区分相机 valid_cams = os.listdir(vid_root_path) #这一行必须加 valid_cams.remove("annotation.json") #给相机排序，方便取出01开头的相机，因为序号小的相机对应的物体更多 valid_cams = natsorted(valid_cams) print(valid_cams) ego_cams = [] exo_cams = [] for vc in valid_cams: if 'aria' in vc: ego_cams.append(vc) else: exo_cams.append(vc) ego = ego_cams[0] exo = exo_cams[0] # print(ego, exo) #ego、exo相机路径 vid_ego_path = os.path.join(vid_root_path, ego) # vid_exo_path = os.path.join(vid_root_path, exo) #setting为ego->exo，所以ego作为第一帧，即visual prompt #取出第一帧ego图像的id #获取帧的索引时，不能简单地通过os.listdir来获取，因为路径下有的图片是没有标注的，需要以注释文件的索引为准 #路径下图片的索引和注释文件中的索引的关系：图片名称里的subsample_idx是包含annotations里的idx的 即有的图片是没有对应的注释的，所以会出现索引报错 ego_frames = natsorted(os.listdir(vid_ego_path)) ego_frames = [int(f.split(".")[0]) for f in ego_frames] #先选出两个摄像机下都出现的物体作为总的物体范围,然后再判断在该摄像机视角下每一帧中出现了哪些物体 #也可能出现objs_both_have为空的情况，这时候就需要更换exo摄像机 objs_both_have = [] for obj in objs: if ego in annotations["masks"][obj].keys() and exo in annotations["masks"][obj].keys(): objs_both_have.append(obj) if len(exo_cams) > 1: for cam in exo_cams[1:]: objs_both_have_tmp = [] for obj in objs: if ego in annotations["masks"][obj].keys() and cam in annotations["masks"][obj].keys(): objs_both_have_tmp.append(obj) if len(objs_both_have_tmp) > len(objs_both_have): exo = cam objs_both_have = objs_both_have_tmp # 如果没有物体范围，跳过本take print("objs_both_have num:", len(objs_both_have)) if len(objs_both_have) == 0: bad_case.append(val_name) continue print(ego, exo) #确定exo的最终相机后，再定义exo的路径 vid_exo_path = os.path.join(vid_root_path, exo) print(f"vid_exo_path:{vid_exo_path}") exo_frames = natsorted(os.listdir(vid_exo_path)) # exo_frames = [int(f.split(".")[0]) for f in exo_frames] # exo_frames是字符串形式 be like ['1' '2' '3'] exo_frames = [f.split(".")[0] for f in exo_frames] # 获取ego注释文件中的所有索引，用于后续和exo的交叉 # 取所有ego obj annotated_frames最长的作为基准帧数 # 后续对exo的操作以基准帧为核心，而不是以物体为核心 # 取ego视角下出现时间最长的物体对应的所有注释帧，作为基准帧 obj_ref = objs_both_have[0] for obj in objs_both_have: if len(list(annotations["masks"][obj_ref][ego].keys())) < len(list(annotations["masks"][obj][ego].keys())): obj_ref = obj ego_anno_frames = natsorted(list(annotations["masks"][obj_ref][ego].keys())) # TODO给frames排个序 frames = natsorted(np.intersect1d(ego_anno_frames, exo_frames)) print(f"frames:{frames}") #查看每一个物体下具体有哪些帧 # ego_anno_frames = natsorted(list(annotations["masks"][objs[0]][ego].keys())) # ego_anno_frames2 = natsorted(list(annotations["masks"][objs[1]][ego].keys())) # print(f"ego_anno_frames3:{ego_anno_frames3}") #TODO测试一下结果是什么样的，默认最好是字符串 #获取ego有注释的第一帧作为参考图像 all_ref_keys = np.asarray( natsorted(annotations["masks"][obj_ref][ego]) ).astype(np.int64) #first_anno_key是ego有注释第一张图片的索引 first_anno_key = str(all_ref_keys[0]) rgb_name = f"{first_anno_key}.jpg" first_frame_img_path = os.path.join(vid_ego_path, rgb_name) first_frame_img_relpath = os.path.relpath(first_frame_img_path, root_path) #实验需改动 # first_frame_img_relpath = "piano_test/aria01_214-1/0.jpg" # first_frame_annotation_img = Image.open(first_frame_annotation_path) # first_frame_annotation = np.array(first_frame_annotation_img) # height, width = first_frame_annotation.shape # 改为通过json文件获取ego mask大小，在我们的脚本中用不上，因为ego和exo大小不一样 # height1, width1 = annotations["masks"][obj_ref][ego][first_anno_key]["size"] #np.unique存储每一帧中的所有像素类别 # unique_instances = np.unique(first_frame_annotation) # unique_instances = unique_instances[unique_instances != 0] #这个列表用于存储第一帧的注释信息 coco_format_annotations = [] #统计每一帧下具体有哪些物体，这里统计的是参考帧ego的 #追踪的物体范围以ego参考帧中的物体为准，因为你输入的mask不可能超过这个范围 obj_list_ego = [] for obj in objs_both_have: if first_anno_key in annotations["masks"][obj][ego].keys(): mask_ego = decode(annotations["masks"][obj][ego][first_anno_key]) area_new = mask_ego.sum().astype(float) if area_new != 0: obj_list_ego.append(obj) print("total obj num in ego", len(obj_list_ego)) if len(obj_list_ego) == 0: bad_case.append(val_name) continue # print(obj_list_ego) # 因为有的exo图像的大小是(960,540)，所以ego mask缩放的大小不能写死 idx_tmp = frames[1] filename_tmp = f"{idx_tmp}.jpg" tmp_path = os.path.join(vid_exo_path, filename_tmp) img_tmp = Image.open(tmp_path) img_tmp = np.array(img_tmp) h_tmp, w_tmp = img_tmp.shape[:2] #处理ego帧中的物体mask obj_list_ego_new = [] for obj in obj_list_ego: #TODO看看segmentation中count和size的顺序影不影响使用 segmentation_tmp = annotations["masks"][obj][ego][first_anno_key] # 可以直接从annotation中取出来 # area可能得decode搞一下 binary_mask = decode(segmentation_tmp) # print("original binary_mask_shape:", binary_mask.shape) #对解码后的mask进行缩放，使得可以匹配ego图像的大小 h,w = binary_mask.shape binary_mask = cv2.resize(binary_mask, (w_tmp, h_tmp), interpolation=cv2.INTER_NEAREST) #这里计算的area是resize后的mask面积 area = binary_mask.sum().astype(float) if area == 0: # obj_list_ego.remove(obj) continue segmentation = encode(np.asfortranarray(binary_mask)) segmentation = { 'counts': segmentation['counts'].decode('ascii'), 'size': segmentation["size"], } obj_list_ego_new.append(obj) coco_format_annotations.append( { 'segmentation': segmentation, 'area': area, 'category_id': float(coco_id_to_cont_id[obj]), } ) if len(obj_list_ego_new) == 0: bad_case.append(val_name) continue #检查每个物体对应哪些摄像机，因为并不是每个物体对应所有的摄像机 # for obj in objs: # cams = list(annotations["masks"][obj].keys()) # print(f"{obj}:{cams}") #开始处理exo相机下的每一帧 #看看索引从1开始还是从0开始 for idx in frames[1:]: filename = f"{idx}.jpg" sample_img_path = os.path.join(vid_exo_path, filename) sample_img_relpath = os.path.relpath(sample_img_path, root_path) #统计每一exo帧下有哪些物体 #有两种方式，第一种是统计该帧下在本take所有物体范围objs中出现的物体，可能会出现Found new target not in the first frame的错误 #第二种方式是统计统计该帧下在参考帧范围obj_list_ego中出现的物体 obj_list_exo = [] for obj in obj_list_ego_new: if idx in annotations["masks"][obj][exo].keys(): mask_exo = decode(annotations["masks"][obj][exo][idx]) area_exo = mask_exo.sum().astype(float) if area_exo != 0: obj_list_exo.append(obj) #检查exo下每一帧的物体数量，也会碰到有的帧一个物体也没有，这种直接跳过 print("total obj num in exo", len(obj_list_exo)) if len(obj_list_exo) == 0: continue height, width = annotations["masks"][obj_list_exo[0]][exo][idx]["size"] # print("original exo mask_shape:" ,height,width) image_info = { 'file_name': sample_img_relpath, 'height': height//4, 'width': width//4, } anns = [] obj_list_exo_new = [] for obj in obj_list_exo: assert obj in obj_list_ego_new, 'Found new target not in the first frame' segmentation_tmp = annotations["masks"][obj][exo][idx] binary_mask = decode(segmentation_tmp) # print("original ego binary_mask_shape", binary_mask.shape) h, w = binary_mask.shape binary_mask = cv2.resize(binary_mask, (w // 4, h // 4), interpolation=cv2.INTER_NEAREST) # print("binary_mask", binary_mask.shape) area = binary_mask.sum().astype(float) if area == 0: continue segmentation = encode(np.asfortranarray(binary_mask)) segmentation = { 'counts': segmentation['counts'].decode('ascii'), 'size': segmentation['size'], } obj_list_exo_new.append(obj) anns.append( { 'segmentation': segmentation, 'area': area, 'category_id': float(coco_id_to_cont_id[obj]), } ) if len(obj_list_exo_new) == 0: continue # 统计本帧下物体对应的id，方便后续根据category_id调整first_frame_anns sample_unique_instances = [float(coco_id_to_cont_id[obj]) for obj in obj_list_exo_new] # 查看每一帧下有哪些物体 print(f"sample_unique_instances in {idx}:{sample_unique_instances}") #deepcopy的目的是，后续要根据本exo帧中物体的数量对参考帧的注释进行调整，防止修改原始注释 first_frame_anns = copy.deepcopy(coco_format_annotations) #考虑本帧物体的数量小于参考帧的情况，仅取出参考帧中本帧有的物体的注释；但是实际情况下，有可能本帧物体的数量会大于参考帧，这时候就需要调整统计本帧下有哪些物体时，总的物体范围 if len(anns) < len(first_frame_anns): first_frame_anns = [ann for ann in first_frame_anns if ann['category_id'] in sample_unique_instances] assert len(anns) == len(first_frame_anns) sample = { 'image': sample_img_relpath, 'image_info': image_info, 'anns': anns, 'first_frame_image': first_frame_img_relpath, 'first_frame_anns': first_frame_anns, 'new_img_id': new_img_id, 'video_name': val_name, } egoexo_dataset.append(sample) new_img_id += 1 print(bad_case) with open(save_path, 'w') as f: json.dump(egoexo_dataset, f) print(f'Save at {save_path}. Total sample: {len(egoexo_dataset)}')