File size: 4,311 Bytes
625a17f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | 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 string
def extract_object_name(text):
parts = text.split("is")
if len(parts) > 1:
return parts[1].strip()
return None
if __name__ == '__main__':
# new_img_id = 0
sent_id = 1
data_new_list = []
json_path = "/data/work-gcp-europe-west4-a/yuqian_fu/Ego/data_segswap/ExoQuery_val_newprompt_all_instruction.json"
with open(json_path, 'r') as f:
datas = json.load(f)
root_path = "/data/work-gcp-europe-west4-a/yuqian_fu/Ego/data_segswap"
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"]
for video_name in tqdm(val_set):
data_thisvideo = []
for data in datas:
if data["video_name"] == video_name:
data_thisvideo.append(data)
if len(data_thisvideo) == 0:
continue
first_frame = data_thisvideo[0]
sample_unique_instances = []
for ann in first_frame['first_frame_anns']:
sample_unique_instances.append(ann['category_id'])
data_new_list.append(first_frame)
for sample in data_thisvideo[1:]:
# sample作为target
if len(first_frame['first_frame_anns']) < len(sample['anns']):
continue
unique_instances = []
#参考帧的物体类别一定要包含当前帧的
for ann in sample['anns']:
unique_instances.append(ann['category_id'])
skip = False
for id in unique_instances:
if id not in sample_unique_instances:
skip = True
break
if skip:
continue
first_frame_anns = copy.deepcopy(first_frame['first_frame_anns'])
if len(sample['anns']) < len(first_frame_anns):
first_frame_anns = [ann for ann in first_frame_anns if ann['category_id'] in unique_instances]
assert len(sample['anns']) == len(first_frame_anns)
skip_text = False #debug
instruct_list = []
for anno in first_frame_anns:
text = anno["text"]
# 提取is之后的句子
raw = extract_object_name(text)
#将raw变小写
if raw == None: #debug
skip_text = True
# print(sample['image'])
break
raw_lower = raw.lower()
# 删除 "green" 并去掉多余的空格
result = raw_lower.replace("green", "").strip()
# 删除所有标点符号
sent = result.translate(str.maketrans('', '', string.punctuation))
tokens = sent.split()
sample_text = {
"tokens": tokens,
"raw": raw,
"sent_id": sent_id,
"sent": sent
}
sent_id += 1
instruct_list.append(sample_text)
if skip_text:
continue
# image_info = {
# "file_name": sample['first_frame_image'],
# "height": 704,
# "width": 704,
# }
data_new = {
'image': sample['image'],
'image_info': sample['image_info'],
'anns': sample['anns'],
'first_frame_image':first_frame['first_frame_image'],
'first_frame_anns': first_frame_anns,
'new_img_id': sample['new_img_id'],
'video_name': sample['video_name'],
"instruction": instruct_list
}
data_new_list.append(data_new)
save_path = "/data/work-gcp-europe-west4-a/yuqian_fu/Ego/data_segswap/Exo_firstframe_VideoTracking.json"
with open(save_path, 'w') as f:
json.dump(data_new_list, f)
print(f'Save at {save_path}. Total sample: {len(data_new_list)}')
|