File size: 2,861 Bytes
032e687 |
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 |
from .ReVOS_Dataset import VideoReVOSDataset
import json
import pickle
class VideoRefYoutubeVOSDataset(VideoReVOSDataset):
def __init__(self,
image_folder,
expression_file,
mask_file,
extra_image_processor=None,
tokenizer=None,
select_number=5,
sampled_frames=10,
offline_processed_text_folder=None,
template_map_fn=None,
max_length=2048,
lazy=True,
repeats=1,
special_tokens=None,
):
super().__init__(
image_folder=image_folder,
expression_file=expression_file,
mask_file=mask_file,
tokenizer=tokenizer,
extra_image_processor=extra_image_processor,
select_number=select_number,
sampled_frames=sampled_frames,
offline_processed_text_folder=offline_processed_text_folder,
template_map_fn=template_map_fn,
max_length=max_length,
lazy=lazy,
repeats=repeats,
special_tokens=special_tokens,
)
def json_file_preprocess(self, expression_file, mask_file):
# prepare expression annotation files
with open(expression_file, 'r') as f:
expression_datas = json.load(f)['videos']
metas = []
anno_count = 0 # serve as anno_id
vid2metaid = {}
for vid_name in expression_datas:
vid_express_data = expression_datas[vid_name]
vid_frames = sorted(vid_express_data['frames'])
vid_len = len(vid_frames)
exp_id_list = sorted(list(vid_express_data['expressions'].keys()))
for exp_id in exp_id_list:
exp_dict = vid_express_data['expressions'][exp_id]
meta = {}
meta['video'] = vid_name
meta['exp'] = exp_dict['exp'] # str
meta['mask_anno_id'] = [str(anno_count), ]
if 'obj_id' in exp_dict.keys():
meta['obj_id'] = exp_dict['obj_id']
else:
meta['obj_id'] = [0, ] # Ref-Youtube-VOS only has one object per expression
meta['anno_id'] = [str(anno_count), ]
anno_count += 1
meta['frames'] = vid_frames
meta['exp_id'] = exp_id
meta['length'] = vid_len
metas.append(meta)
if vid_name not in vid2metaid.keys():
vid2metaid[vid_name] = []
vid2metaid[vid_name].append(len(metas) - 1)
# process mask annotation files
with open(mask_file, 'rb') as f:
mask_dict = pickle.load(f)
return vid2metaid, metas, mask_dict
|