|
|
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
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
root_path = '/data/work2-gcp-europe-west4-a/yuqian_fu/Ego/data_segswap'
|
|
|
|
|
|
save_path = os.path.join(root_path, 'exoego_nullmask_eval.json')
|
|
|
|
|
|
|
|
|
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"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
new_img_id = 0
|
|
|
|
|
|
|
|
|
egoexo_dataset = []
|
|
|
|
|
|
'''
|
|
|
build_DAVIS.py的代码逻辑是先处理每个视频的第一帧，第一帧中的unique_instances、高宽等信息用于该视频下后续的每一帧。
|
|
|
注意，unique_instances代表的是第一帧下像素的所有类别信息，如果该视频下后续的帧中有像素的类别不在unique_instances中，会报错
|
|
|
'''
|
|
|
|
|
|
|
|
|
|
|
|
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)
|
|
|
|
|
|
|
|
|
objs = natsorted(list(annotations["masks"].keys()))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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")
|
|
|
|
|
|
|
|
|
valid_cams = natsorted(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]
|
|
|
|
|
|
|
|
|
|
|
|
vid_ego_path = os.path.join(vid_root_path, ego)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
ego_frames = natsorted(os.listdir(vid_ego_path))
|
|
|
ego_frames = [f.split(".")[0] for f in ego_frames]
|
|
|
|
|
|
|
|
|
|
|
|
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
|
|
|
|
|
|
|
|
|
if len(objs_both_have) == 0:
|
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
|
|
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 = [f.split(".")[0] for f in exo_frames]
|
|
|
|
|
|
|
|
|
|
|
|
obj_ref = objs_both_have[0]
|
|
|
for obj in objs_both_have:
|
|
|
if len(list(annotations["masks"][obj_ref][exo].keys())) < len(list(annotations["masks"][obj][exo].keys())):
|
|
|
obj_ref = obj
|
|
|
exo_anno_frames = natsorted(list(annotations["masks"][obj_ref][exo].keys()))
|
|
|
|
|
|
frames = natsorted(np.intersect1d(exo_anno_frames, ego_frames))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
for idx in frames:
|
|
|
coco_format_annotations = []
|
|
|
filename = f"{idx}.jpg"
|
|
|
sample_img_path = os.path.join(vid_ego_path, filename)
|
|
|
sample_img_relpath = os.path.relpath(sample_img_path, root_path)
|
|
|
|
|
|
first_frame_img_path = os.path.join(vid_exo_path, filename)
|
|
|
first_frame_img_relpath = os.path.relpath(first_frame_img_path, root_path)
|
|
|
|
|
|
obj_list_exo = []
|
|
|
for obj in objs_both_have:
|
|
|
if idx in annotations["masks"][obj][exo].keys():
|
|
|
mask_exo = decode(annotations["masks"][obj][exo][idx])
|
|
|
area_new = mask_exo.sum().astype(float)
|
|
|
if area_new != 0:
|
|
|
obj_list_exo.append(obj)
|
|
|
|
|
|
if len(obj_list_exo) == 0:
|
|
|
continue
|
|
|
|
|
|
obj_list_exo_new = []
|
|
|
for obj in obj_list_exo:
|
|
|
|
|
|
segmentation_tmp = annotations["masks"][obj][exo][idx]
|
|
|
binary_mask = decode(segmentation_tmp)
|
|
|
|
|
|
|
|
|
|
|
|
h, w = binary_mask.shape
|
|
|
binary_mask = cv2.resize(binary_mask, (w // 4, h // 4), interpolation=cv2.INTER_NEAREST)
|
|
|
|
|
|
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)
|
|
|
coco_format_annotations.append(
|
|
|
{
|
|
|
'segmentation': segmentation,
|
|
|
'area': area,
|
|
|
'category_id': float(coco_id_to_cont_id[obj]),
|
|
|
}
|
|
|
)
|
|
|
if len(obj_list_exo_new) == 0:
|
|
|
continue
|
|
|
|
|
|
obj_list_ego = []
|
|
|
for obj in obj_list_exo_new:
|
|
|
if idx in annotations["masks"][obj][ego].keys():
|
|
|
mask_ego = decode(annotations["masks"][obj][ego][idx])
|
|
|
area_ego = mask_ego.sum().astype(float)
|
|
|
if area_ego != 0:
|
|
|
obj_list_ego.append(obj)
|
|
|
|
|
|
|
|
|
|
|
|
if len(obj_list_ego) == 0:
|
|
|
continue
|
|
|
|
|
|
height, width = annotations["masks"][obj_list_ego[0]][ego][idx]["size"]
|
|
|
|
|
|
image_info = {
|
|
|
'file_name': sample_img_relpath,
|
|
|
'height': height // 2,
|
|
|
'width': width // 2,
|
|
|
}
|
|
|
|
|
|
anns = []
|
|
|
|
|
|
|
|
|
obj_list_ego_new = []
|
|
|
|
|
|
for obj in obj_list_ego:
|
|
|
assert obj in obj_list_exo_new, 'Found new target not in the first frame'
|
|
|
segmentation_tmp = annotations["masks"][obj][ego][idx]
|
|
|
binary_mask = decode(segmentation_tmp)
|
|
|
|
|
|
h, w = binary_mask.shape
|
|
|
binary_mask = cv2.resize(binary_mask, (w // 2, h // 2), interpolation=cv2.INTER_NEAREST)
|
|
|
|
|
|
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_ego_new.append(obj)
|
|
|
anns.append(
|
|
|
{
|
|
|
'segmentation': segmentation,
|
|
|
'area': area,
|
|
|
'category_id': float(coco_id_to_cont_id[obj]),
|
|
|
}
|
|
|
)
|
|
|
|
|
|
|
|
|
if len(obj_list_ego_new) == 0:
|
|
|
continue
|
|
|
|
|
|
|
|
|
sample_unique_instances = [float(coco_id_to_cont_id[obj]) for obj in obj_list_ego_new]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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
|
|
|
|
|
|
with open(save_path, 'w') as f:
|
|
|
json.dump(egoexo_dataset, f)
|
|
|
print(f'Save at {save_path}. Total sample: {len(egoexo_dataset)}')
|
|
|
|
