Spaces:
Sleeping
Sleeping
| # Copyright (c) Facebook, Inc. All Rights Reserved | |
| import numpy as np | |
| import os | |
| import torch | |
| class Processor(object): | |
| """ | |
| A generic processor for video (codec, feature etc.) and text. | |
| """ | |
| def __call__(self, **kwargs): | |
| raise NotImplementedError | |
| class MetaProcessor(Processor): | |
| """ | |
| A meta processor is expected to load the metadata of a dataset: | |
| (e.g., video_ids, or captions). | |
| You must implement the `__getitem__` (meta datasets are rather diverse.). | |
| """ | |
| def __init__(self, config): | |
| self.split = config.split | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, idx): | |
| raise NotImplementedError | |
| def _get_split_path(self, config): | |
| splits = { | |
| "train": config.train_path, | |
| "valid": config.val_path, | |
| "test": config.test_path, | |
| } | |
| if config.split is not None: | |
| return splits[config.split] | |
| return config.train_path | |
| class TextProcessor(Processor): | |
| """ | |
| A generic Text processor: rename this as `withTokenizer`. | |
| tokenize a string of text on-the-fly. | |
| Warning: mostly used for end tasks. | |
| (on-the-fly tokenization is slow for how2.) | |
| TODO(huxu): move this class as a subclass. | |
| """ | |
| def __init__(self, config): | |
| self.bert_name = str(config.bert_name) | |
| self.use_fast = config.use_fast | |
| from transformers import AutoTokenizer | |
| self.tokenizer = AutoTokenizer.from_pretrained( | |
| self.bert_name, use_fast=self.use_fast | |
| ) | |
| def __call__(self, text_id): | |
| caption = self.tokenizer(text_id, add_special_tokens=False) | |
| return caption["input_ids"] | |
| class VideoProcessor(Processor): | |
| """ | |
| A generic video processor: load a numpy video tokens by default. | |
| """ | |
| def __init__(self, config): | |
| self.vfeat_dir = config.vfeat_dir | |
| def __call__(self, video_fn): | |
| if isinstance(video_fn, tuple): | |
| video_fn = video_fn[0] | |
| assert isinstance(video_fn, str) | |
| video_fn = os.path.join(self.vfeat_dir, video_fn + ".npy") | |
| feat = np.load(video_fn) | |
| return feat | |
| class Aligner(object): | |
| """ | |
| An alignprocessor align video and text and output a dict of tensors (for a model). | |
| """ | |
| def __init__(self, config): | |
| """__init__ needs to be light weight for more workers/threads.""" | |
| self.split = config.split | |
| self.max_video_len = config.max_video_len | |
| self.max_len = config.max_len | |
| from transformers import AutoTokenizer | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| str(config.bert_name), use_fast=config.use_fast | |
| ) | |
| self.cls_token_id = tokenizer.cls_token_id | |
| self.sep_token_id = tokenizer.sep_token_id | |
| self.pad_token_id = tokenizer.pad_token_id | |
| self.mask_token_id = tokenizer.mask_token_id | |
| def __call__(self, video_id, video_feature, text_feature): | |
| raise NotImplementedError | |
| def _build_video_seq(self, video_feature, video_clips=None): | |
| """ | |
| `video_feature`: available video tokens. | |
| `video_clips`: video clip sequence to build. | |
| """ | |
| if not isinstance(video_feature, np.ndarray): | |
| raise ValueError( | |
| "unsupported type of video_feature", type(video_feature) | |
| ) | |
| if video_clips is None: | |
| # this is borrowed from DSAligner | |
| video_start = 0 | |
| video_end = min(len(video_feature), self.max_video_len) | |
| # the whole sequence is a single clip. | |
| video_clips = {"start": [video_start], "end": [video_end]} | |
| vfeats = np.zeros( | |
| (self.max_video_len, video_feature.shape[1]), dtype=np.float32 | |
| ) | |
| vmasks = torch.zeros((self.max_video_len,), dtype=torch.bool) | |
| video_len = 0 | |
| for start, end in zip(video_clips["start"], video_clips["end"]): | |
| clip_len = min(self.max_video_len - video_len, (end - start)) | |
| if clip_len > 0: | |
| vfeats[video_len: video_len + clip_len] = video_feature[ | |
| start: start + clip_len | |
| ] | |
| vmasks[video_len: video_len + clip_len] = 1 | |
| video_len += clip_len | |
| vfeats = torch.from_numpy(vfeats) | |
| return vfeats, vmasks | |
| def _build_text_seq(self, text_feature, text_clip_indexs=None): | |
| """ | |
| `text_feature`: all available clips. | |
| `text_clip_indexes`: clip sequence to build. | |
| """ | |
| if text_clip_indexs is None: | |
| text_clip_indexs = [0] | |
| full_caps = [] | |
| if isinstance(text_feature, dict): | |
| for clip_idx in text_clip_indexs: | |
| full_caps.extend(text_feature["cap"][clip_idx]) | |
| else: | |
| full_caps = text_feature | |
| max_text_len = self.max_len - self.max_video_len - 3 | |
| full_caps = full_caps[:max_text_len] | |
| full_caps = ( | |
| [self.cls_token_id, self.sep_token_id] + full_caps + [self.sep_token_id] | |
| ) | |
| text_pad_len = self.max_len - len(full_caps) - self.max_video_len | |
| padded_full_caps = full_caps + [self.pad_token_id] * text_pad_len | |
| caps = torch.LongTensor(padded_full_caps) | |
| cmasks = torch.zeros((len(padded_full_caps),), dtype=torch.bool) | |
| cmasks[: len(full_caps)] = 1 | |
| return caps, cmasks | |
| def batch_post_processing(self, batch, video_feature): | |
| return batch | |
| class MMAttentionMask2DProcessor(Processor): | |
| """text generation requires 2d mask | |
| that is harder to generate by GPU at this stage.""" | |
| def __call__(self, vmask, cmask, mtype): | |
| if mtype == "textgen": | |
| return self._build_textgeneration_mask(vmask, cmask) | |
| elif mtype == "videogen": | |
| return self._build_videogeneration_mask(vmask, cmask) | |
| else: | |
| return self._build_mm_mask(vmask, cmask) | |
| def _build_mm_mask(self, vmask, cmask): | |
| mask_1d = torch.cat([cmask[:1], vmask, cmask[1:]], dim=0) | |
| return mask_1d[None, :].repeat(mask_1d.size(0), 1) | |
| def _build_videogeneration_mask(self, vmask, cmask): | |
| # cls_mask is only about text otherwise it will leak generation. | |
| cls_text_mask = torch.cat([ | |
| # [CLS] | |
| torch.ones( | |
| (1,), dtype=torch.bool, device=cmask.device), | |
| # video tokens and [SEP] for video. | |
| torch.zeros( | |
| (vmask.size(0) + 1,), dtype=torch.bool, device=cmask.device), | |
| cmask[2:] | |
| ], dim=0) | |
| # concat horizontially. | |
| video_len = int(vmask.sum()) | |
| video_masks = torch.cat([ | |
| # [CLS] | |
| torch.ones( | |
| (video_len, 1), dtype=torch.bool, device=cmask.device | |
| ), | |
| torch.tril( | |
| torch.ones( | |
| (video_len, video_len), | |
| dtype=torch.bool, device=cmask.device)), | |
| # video_padding | |
| torch.zeros( | |
| (video_len, vmask.size(0) - video_len), | |
| dtype=torch.bool, device=cmask.device | |
| ), | |
| # [SEP] for video (unused). | |
| torch.zeros( | |
| (video_len, 1), dtype=torch.bool, device=cmask.device | |
| ), | |
| cmask[2:].unsqueeze(0).repeat(video_len, 1) | |
| ], dim=1) | |
| text_masks = cls_text_mask[None, :].repeat( | |
| cmask.size(0) - 2, 1) | |
| video_padding_masks = cls_text_mask[None, :].repeat( | |
| vmask.size(0) - video_len, 1) | |
| return torch.cat([ | |
| cls_text_mask[None, :], | |
| video_masks, | |
| video_padding_masks, | |
| torch.cat([cmask[:1], vmask, cmask[1:]], dim=0)[None,:], | |
| text_masks | |
| ], dim=0) | |
| def _build_textgeneration_mask(self, vmask, cmask): | |
| # cls_mask is only about video otherwise it will leak generation. | |
| cls_video_mask = torch.cat([ | |
| # [CLS] | |
| torch.ones( | |
| (1,), dtype=torch.bool, device=cmask.device), | |
| vmask, | |
| # [SEP] | |
| torch.ones((1,), dtype=torch.bool, device=cmask.device), | |
| torch.zeros( | |
| (cmask.size(0)-2,), dtype=torch.bool, device=cmask.device) | |
| ], dim=0) | |
| # concat horizontially. | |
| text_len = int(cmask[2:].sum()) | |
| text_masks = torch.cat([ | |
| # [CLS] | |
| torch.ones( | |
| (text_len, 1), dtype=torch.bool, device=cmask.device | |
| ), | |
| vmask.unsqueeze(0).repeat(text_len, 1), | |
| # [SEP] for video. | |
| torch.ones( | |
| (text_len, 1), dtype=torch.bool, device=cmask.device | |
| ), | |
| torch.tril( | |
| torch.ones( | |
| (text_len, text_len), | |
| dtype=torch.bool, device=cmask.device)), | |
| # padding. | |
| torch.zeros( | |
| (text_len, cmask.size(0) - text_len - 2), | |
| dtype=torch.bool, device=cmask.device | |
| ) | |
| ], dim=1) | |
| cls_video_masks = cls_video_mask[None, :].repeat( | |
| vmask.size(0) + 2, 1) | |
| text_padding_masks = cls_video_mask[None, :].repeat( | |
| cmask.size(0) - text_len - 2, 1) | |
| return torch.cat([ | |
| cls_video_masks, text_masks, text_padding_masks], dim=0) | |