From 5debc88edf986e99a3a3c209077519e884a3c1fb Mon Sep 17 00:00:00 2001 From: "gongbiao.gb" Date: Wed, 11 Jan 2023 22:22:12 +0800 Subject: [PATCH] add vop retrieval Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11302991 --- modelscope/metainfo.py | 2 + modelscope/models/cv/__init__.py | 3 +- .../models/cv/vop_retrieval/__init__.py | 28 ++ .../models/cv/vop_retrieval/backbone.py | 354 ++++++++++++++++ .../models/cv/vop_retrieval/basic_utils.py | 170 ++++++++ modelscope/models/cv/vop_retrieval/model.py | 378 ++++++++++++++++++ .../cv/vop_retrieval/tokenization_clip.py | 159 ++++++++ modelscope/pipelines/cv/__init__.py | 2 + .../pipelines/cv/vop_retrieval_pipeline.py | 122 ++++++ modelscope/utils/constant.py | 1 + pose_keypoint.jpg | 3 + requirements/cv.txt | 2 + tests/pipelines/test_vop_retrieval.py | 36 ++ 13 files changed, 1259 insertions(+), 1 deletion(-) create mode 100644 modelscope/models/cv/vop_retrieval/__init__.py create mode 100644 modelscope/models/cv/vop_retrieval/backbone.py create mode 100644 modelscope/models/cv/vop_retrieval/basic_utils.py create mode 100644 modelscope/models/cv/vop_retrieval/model.py create mode 100644 modelscope/models/cv/vop_retrieval/tokenization_clip.py create mode 100644 modelscope/pipelines/cv/vop_retrieval_pipeline.py create mode 100644 pose_keypoint.jpg create mode 100644 tests/pipelines/test_vop_retrieval.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 45bc2fad..e3b076cb 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -73,6 +73,7 @@ class Models(object): real_basicvsr = 'real-basicvsr' rcp_sceneflow_estimation = 'rcp-sceneflow-estimation' image_casmvs_depth_estimation = 'image-casmvs-depth-estimation' + vop_retrieval_model = 'vop-retrieval-model' ddcolor = 'ddcolor' image_face_fusion = 'image-face-fusion' @@ -287,6 +288,7 @@ class Pipelines(object): video_super_resolution = 'realbasicvsr-video-super-resolution' pointcloud_sceneflow_estimation = 'pointcloud-sceneflow-estimation' image_multi_view_depth_estimation = 'image-multi-view-depth-estimation' + vop_retrieval = 'vop-video-text-retrieval' ddcolor_image_colorization = 'ddcolor-image-colorization' image_face_fusion = 'image-face-fusion' diff --git a/modelscope/models/cv/__init__.py b/modelscope/models/cv/__init__.py index 7e571ad6..ead5cec1 100644 --- a/modelscope/models/cv/__init__.py +++ b/modelscope/models/cv/__init__.py @@ -18,6 +18,7 @@ from . import (action_recognition, animal_recognition, body_2d_keypoints, shop_segmentation, super_resolution, video_frame_interpolation, video_object_segmentation, video_single_object_tracking, video_stabilization, video_summarization, - video_super_resolution, virual_tryon, vision_middleware) + video_super_resolution, virual_tryon, vision_middleware, + vop_retrieval) # yapf: enable diff --git a/modelscope/models/cv/vop_retrieval/__init__.py b/modelscope/models/cv/vop_retrieval/__init__.py new file mode 100644 index 00000000..5b3e762c --- /dev/null +++ b/modelscope/models/cv/vop_retrieval/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .basic_utils import set_seed, get_state_dict, load_data, init_transform_dict, load_frames_from_video + from .model import VoP + from .tokenization_clip import LengthAdaptiveTokenizer +else: + _import_structure = { + 'basic_utils': [ + 'set_seed', 'get_state_dict', 'load_data', 'init_transform_dict', + 'load_frames_from_video' + ], + 'model': ['VoP'], + 'tokenization_clip': ['LengthAdaptiveTokenizer'] + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/cv/vop_retrieval/backbone.py b/modelscope/models/cv/vop_retrieval/backbone.py new file mode 100644 index 00000000..a2b26e07 --- /dev/null +++ b/modelscope/models/cv/vop_retrieval/backbone.py @@ -0,0 +1,354 @@ +# The implementation here is modified based on HuggingFace, originally Apache 2.0 License +# and publicly avaialbe at https://github.com/huggingface/transformers +# Copyright 2018 The HuggingFace Inc. team. +# Copyright 2021-2022 The Alibaba Fundamental Vision Team Authors. All rights reserved. + +import hashlib +import os +import urllib +import warnings +from collections import OrderedDict +from typing import Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from tqdm import tqdm + +from modelscope.models.base.base_torch_model import TorchModel + + +class LayerNorm(nn.LayerNorm): + + def forward(self, x: torch.Tensor): + orig_type = x.dtype + ret = super().forward(x.type(torch.float32)) + return ret.type(orig_type) + + +class QuickGELU(TorchModel): + + def forward(self, x: torch.Tensor): + return x * torch.sigmoid(1.702 * x) + + +class ResidualAttentionBlock(TorchModel): + + def __init__(self, + d_model: int, + n_head: int, + attn_mask: torch.Tensor = None): + super().__init__() + self.attn = nn.MultiheadAttention(d_model, n_head) + self.ln_1 = LayerNorm(d_model) + self.mlp = nn.Sequential( + OrderedDict([('c_fc', nn.Linear(d_model, d_model * 4)), + ('gelu', QuickGELU()), + ('c_proj', nn.Linear(d_model * 4, d_model))])) + self.ln_2 = LayerNorm(d_model) + self.attn_mask = attn_mask + + def attention(self, x: torch.Tensor): + self.attn_mask = self.attn_mask.to( + dtype=x.dtype, + device=x.device) if self.attn_mask is not None else None + return self.attn( + x, x, x, need_weights=False, attn_mask=self.attn_mask)[0] + + def forward(self, x: torch.Tensor): + x = x + self.attention(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + +class Transformer(TorchModel): + + def __init__(self, + width: int, + layers: int, + heads: int, + attn_mask: torch.Tensor = None): + super().__init__() + self.width = width + self.layers = layers + self.resblocks = nn.Sequential(*[ + ResidualAttentionBlock(width, heads, attn_mask) + for _ in range(layers) + ]) + + def forward(self, x: torch.Tensor): + return self.resblocks(x) + + +class VisualTransformer(TorchModel): + + def __init__(self, input_resolution: int, patch_size: int, width: int, + layers: int, heads: int, output_dim: int): + super().__init__() + self.input_resolution = input_resolution + self.output_dim = output_dim + self.conv1 = nn.Conv2d( + in_channels=3, + out_channels=width, + kernel_size=patch_size, + stride=patch_size, + bias=False) + + scale = width**-0.5 + self.class_embedding = nn.Parameter(scale * torch.randn(width)) + self.positional_embedding = nn.Parameter(scale * torch.randn( + (input_resolution // patch_size)**2 + 1, width)) + self.ln_pre = LayerNorm(width) + + self.transformer = Transformer(width, layers, heads) + + self.ln_post = LayerNorm(width) + self.proj = nn.Parameter(scale * torch.randn(width, output_dim)) + + def forward(self, x: torch.Tensor): + x = self.conv1(x) + x = x.reshape(x.shape[0], x.shape[1], -1) + x = x.permute(0, 2, 1) + x_1 = self.class_embedding.to(x.dtype) + x_2 = torch.zeros( + x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device) + x_1 = x_1 + x_2 + x = torch.cat([x_1, x], dim=1) + x = x + self.positional_embedding.to(x.dtype) + x = self.ln_pre(x) + + x = x.permute(1, 0, 2) + x = self.transformer(x) + x = x.permute(1, 0, 2) + + x = self.ln_post(x[:, 0, :]) + + if self.proj is not None: + x = x @ self.proj + + return x + + +class CLIP(TorchModel): + + def __init__(self, embed_dim: int, image_resolution: int, + vision_layers: Union[Tuple[int, int, int, int], int], + vision_width: int, vision_patch_size: int, + context_length: int, vocab_size: int, transformer_width: int, + transformer_heads: int, transformer_layers: int): + super().__init__() + + self.context_length = context_length + + vision_heads = vision_width // 64 + self.visual = VisualTransformer( + input_resolution=image_resolution, + patch_size=vision_patch_size, + width=vision_width, + layers=vision_layers, + heads=vision_heads, + output_dim=embed_dim) + + self.transformer = Transformer( + width=transformer_width, + layers=transformer_layers, + heads=transformer_heads, + attn_mask=self.build_attention_mask()) + + self.vocab_size = vocab_size + self.token_embedding = nn.Embedding(vocab_size, transformer_width) + self.positional_embedding = nn.Parameter( + torch.empty(self.context_length, transformer_width)) + self.ln_final = LayerNorm(transformer_width) + + self.text_projection = nn.Parameter( + torch.empty(transformer_width, embed_dim)) + self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) + + self.initialize_parameters() + + def initialize_parameters(self): + nn.init.normal_(self.token_embedding.weight, std=0.02) + nn.init.normal_(self.positional_embedding, std=0.01) + + proj_std = (self.transformer.width**-0.5) * ( + (2 * self.transformer.layers)**-0.5) + attn_std = self.transformer.width**-0.5 + fc_std = (2 * self.transformer.width)**-0.5 + for block in self.transformer.resblocks: + nn.init.normal_(block.attn.in_proj_weight, std=attn_std) + nn.init.normal_(block.attn.out_proj.weight, std=proj_std) + nn.init.normal_(block.mlp.c_fc.weight, std=fc_std) + nn.init.normal_(block.mlp.c_proj.weight, std=proj_std) + + if self.text_projection is not None: + nn.init.normal_( + self.text_projection, std=self.transformer.width**-0.5) + + def build_attention_mask(self): + mask = torch.empty(self.context_length, self.context_length) + mask.fill_(float('-inf')) + mask.triu_(1) + return mask + + @property + def dtype(self): + return self.visual.conv1.weight.dtype + + def encode_image(self, image): + return self.visual(image.type(self.dtype)) + + def encode_text(self, text, return_all_tokens=False): + x = self.token_embedding(text).type(self.dtype) + + x = x + self.positional_embedding.type(self.dtype) + x = x.permute(1, 0, 2) + x = self.transformer(x) + x = x.permute(1, 0, 2) + x = self.ln_final(x).type(self.dtype) + + if return_all_tokens: + return x @ self.text_projection + x = x[torch.arange(x.shape[0]), + text.argmax(dim=-1)] @ self.text_projection + + return x + + def forward(self, image, text): + image_features = self.encode_image(image) + text_features = self.encode_text(text) + image_features = image_features / image_features.norm( + dim=-1, keepdim=True) + text_features = text_features / text_features.norm( + dim=-1, keepdim=True) + logit_scale = self.logit_scale.exp() + logits_per_image = logit_scale * image_features @ text_features.t() + logits_per_text = logit_scale * text_features @ image_features.t() + return logits_per_image, logits_per_text + + +def build_model(state_dict: dict): + vit = 'visual.proj' in state_dict + + if vit: + vision_width = state_dict['visual.conv1.weight'].shape[0] + vision_layers = len([ + k for k in state_dict.keys() + if k.startswith('visual.') and k.endswith('.attn.in_proj_weight') + ]) + vision_patch_size = state_dict['visual.conv1.weight'].shape[-1] + grid_size = round( + (state_dict['visual.positional_embedding'].shape[0] - 1)**0.5) + image_resolution = vision_patch_size * grid_size + else: + counts: list = [ + len( + set( + k.split('.')[2] for k in state_dict + if k.startswith(f'visual.layer{b}'))) + for b in [1, 2, 3, 4] + ] + vision_layers = tuple(counts) + vision_width = state_dict['visual.layer1.0.conv1.weight'].shape[0] + output_width = round( + (state_dict['visual.attnpool.positional_embedding'].shape[0] + - 1)**0.5) + vision_patch_size = None + assert output_width**2 + 1 == state_dict[ + 'visual.attnpool.positional_embedding'].shape[0] + image_resolution = output_width * 32 + + embed_dim = state_dict['text_projection'].shape[1] + context_length = state_dict['positional_embedding'].shape[0] + vocab_size = state_dict['token_embedding.weight'].shape[0] + transformer_width = state_dict['ln_final.weight'].shape[0] + transformer_heads = transformer_width // 64 + transformer_layers = len( + set( + k.split('.')[2] for k in state_dict + if k.startswith('transformer.resblocks'))) + + model = CLIP(embed_dim, image_resolution, vision_layers, vision_width, + vision_patch_size, context_length, vocab_size, + transformer_width, transformer_heads, transformer_layers) + + for key in ['input_resolution', 'context_length', 'vocab_size']: + if key in state_dict: + del state_dict[key] + + model.load_state_dict(state_dict) + return model.eval() + + +def load_clip(name: str, + device: Union[str, torch.device] = 'cuda' + if torch.cuda.is_available() else 'cpu', + jit=True): + jit = False + model_path = name + try: + model = torch.jit.load( + model_path, map_location=device if jit else 'cpu').eval() + state_dict = None + except RuntimeError: + if jit: + warnings.warn( + f'File {model_path} is not a JIT archive. Loading as a state dict instead' + ) + jit = False + state_dict = torch.load(model_path, map_location='cpu') + + if not jit: + model = build_model(state_dict or model.state_dict()).to(device) + if str(device) == 'cpu': + model.float() + return model + + device_holder = torch.jit.trace( + lambda: torch.ones([]).to(torch.device(device)), example_inputs=[]) + device_node = [ + n for n in device_holder.graph.findAllNodes('prim::Constant') + if 'Device' in repr(n) + ][-1] + + def patch_device(module): + graphs = [module.graph] if hasattr(module, 'graph') else [] + if hasattr(module, 'forward1'): + graphs.append(module.forward1.graph) + + for graph in graphs: + for node in graph.findAllNodes('prim::Constant'): + if 'value' in node.attributeNames() and str( + node['value']).startswith('cuda'): + node.copyAttributes(device_node) + + model.apply(patch_device) + patch_device(model.encode_image) + patch_device(model.encode_text) + + if str(device) == 'cpu': + float_holder = torch.jit.trace( + lambda: torch.ones([]).float(), example_inputs=[]) + float_input = list(float_holder.graph.findNode('aten::to').inputs())[1] + float_node = float_input.node() + + def patch_float(module): + graphs = [module.graph] if hasattr(module, 'graph') else [] + if hasattr(module, 'forward1'): + graphs.append(module.forward1.graph) + + for graph in graphs: + for node in graph.findAllNodes('aten::to'): + inputs = list(node.inputs()) + for i in [1, 2]: + if inputs[i].node()['value'] == 5: + inputs[i].node().copyAttributes(float_node) + + model.apply(patch_float) + patch_float(model.encode_image) + patch_float(model.encode_text) + + model.float() + + return model diff --git a/modelscope/models/cv/vop_retrieval/basic_utils.py b/modelscope/models/cv/vop_retrieval/basic_utils.py new file mode 100644 index 00000000..c38482fc --- /dev/null +++ b/modelscope/models/cv/vop_retrieval/basic_utils.py @@ -0,0 +1,170 @@ +# Copyright 2021-2022 The Alibaba Fundamental Vision Team Authors. All rights reserved. + +import os +import pickle +import random +import shutil +import zipfile +from collections import OrderedDict + +import cv2 +import numpy as np +import torch +import ujson as json +from PIL import Image +from torchvision import transforms + + +def init_transform_dict(input_res=224): + """ + The implementation of transforms functions. + The default image resolution is 224. + The normalize parameter follows the mainstream setting. + """ + tsfm_dict = { + 'clip_test': + transforms.Compose([ + transforms.Resize(input_res, interpolation=Image.BICUBIC), + transforms.CenterCrop(input_res), + transforms.Normalize((0.48145466, 0.4578275, 0.40821073), + (0.26862954, 0.26130258, 0.27577711)), + ]), + 'clip_train': + transforms.Compose([ + transforms.RandomResizedCrop(input_res, scale=(0.5, 1.0)), + transforms.RandomHorizontalFlip(), + transforms.ColorJitter(brightness=0, saturation=0, hue=0), + transforms.Normalize((0.48145466, 0.4578275, 0.40821073), + (0.26862954, 0.26130258, 0.27577711)), + ]) + } + return tsfm_dict + + +def load_data(feature_path, mydevice): + """ + Loading dataset from 'feature_path' as a retrieval docs. + The default dataset is MSRVTT-9K. + + Args: + feature_path: 'VoP_msrvtt9k_features.pkl' + mydevice: device(type='cuda', index=0) + + Returns: + [text_embeds, vid_embeds_pooled, vid_ids, texts] + """ + feature_content = torch.load(feature_path) + text_embeds = feature_content['text_embeds'].to(device=mydevice) + vid_embeds_pooled = feature_content['vid_embeds'].to(device=mydevice) + vid_ids = feature_content['vid_ids'] + texts = feature_content['texts'] + return [text_embeds, vid_embeds_pooled, vid_ids, texts] + + +def load_json(filename): + """ + Load json files. + """ + with open(filename, 'r') as f: + return json.load(f) + + +def set_seed(seed): + """ + Set random seed. + """ + if seed >= 0: + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed_all(seed) + random.seed(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def get_state_dict(checkpoint_path): + """ + Load pre-train parameters for VoP. + """ + checkpoint = torch.load(checkpoint_path) + state_dict = checkpoint['state_dict'] + + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + new_state_dict[k.replace('module.', '')] = v + state_dict = new_state_dict + + return state_dict + + +def get_valid_frames(cap, num_frames, vlen, sample='rand'): + """ + Get indexes of sampled frames. + + Args: + cap: cv2.VideoCapture + num_frames: int - number of frames to sample + vlen: video length, int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 325 + sample: 'rand' | 'uniform' how to sample + + Returns: + frames: torch.tensor of stacked sampled video frames + of dim (num_frames, C, H, W) + frame_idxs: list(int) indices of where the frames where sampled + """ + acc_samples = min(num_frames, vlen) + intervals = np.linspace( + start=0, stop=vlen, num=acc_samples + 1).astype(int) + ranges = [] + for idx, interv in enumerate(intervals[:-1]): + ranges.append((interv, intervals[idx + 1] - 1)) + if sample == 'rand': + frame_idxs = [random.choice(range(x[0], x[1])) for x in ranges] + else: + frame_idxs = [(x[0] + x[1]) // 2 for x in ranges] + + frames = [] + for index in frame_idxs: + cap.set(cv2.CAP_PROP_POS_FRAMES, index) + ret, frame = cap.read() + if not ret: + n_tries = 5 + for _ in range(n_tries): + ret, frame = cap.read() + if ret: + break + if not ret: + return None, None + + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = torch.from_numpy(frame) + frame = frame.permute(2, 0, 1) + frames.append(frame) + + while len(frames) < num_frames: + frames.append(frames[-1].clone()) + + return frames, frame_idxs + + +def load_frames_from_video(video_path, num_frames, sample='rand'): + """ + Get indexes of sampled frames. + + Args: + video_path: the local video path + num_frames: Frame number, 12 frames for each video + sample: 'rand' | 'uniform' how to sample + + Returns: + frames: torch.tensor of stacked sampled video frames + of dim (num_frames, C, H, W) + frame_idxs: list(int) indices of where the frames where sampled + """ + cap = cv2.VideoCapture(video_path) + assert (cap.isOpened()), video_path + vlen = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + frames, frame_idxs = get_valid_frames(cap, num_frames, vlen, sample) + frames = torch.stack(frames).float() / 255 + cap.release() + return frames, frame_idxs diff --git a/modelscope/models/cv/vop_retrieval/model.py b/modelscope/models/cv/vop_retrieval/model.py new file mode 100644 index 00000000..d89fcb7f --- /dev/null +++ b/modelscope/models/cv/vop_retrieval/model.py @@ -0,0 +1,378 @@ +# Copyright 2021-2022 The Alibaba Fundamental Vision Team Authors. All rights reserved. + +import os +import os.path as osp + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from modelscope.metainfo import Models +from modelscope.models.base.base_torch_model import TorchModel +from modelscope.models.builder import MODELS +from modelscope.utils.config import Config +from modelscope.utils.constant import ModelFile, Tasks +from .backbone import load_clip +from .basic_utils import get_state_dict, set_seed + + +@MODELS.register_module( + Tasks.vop_retrieval, module_name=Models.vop_retrieval_model) +class VoP(TorchModel): + """ + The implementation of 'VoP: Text-Video Co-operative Prompt Tuning for Cross-Modal Retrieval'. + This model is dynamically initialized with the following parts: + - clip: the upstream pre-trained backbone model (CLIP in this code) + - pool_frames: the frames pooling method + - visual_prompt_learner: visual prompt + - ImageEncoder: get image encoder + - TextPromptLearner: text prompt + - TextEncoder: get text encoder + """ + + def __init__(self, model_dir: str, *args, **kwargs): + """ + Initialize a VoP Model + + Args: + model_dir: model id or path, + """ + super(VoP, self).__init__() + model_path = osp.join(model_dir, 'VoP_msrvtt9k.pth') + clip_arch = osp.join(model_dir, 'ViT-B-32.pt') + config_path = osp.join(model_dir, ModelFile.CONFIGURATION) + + self.config = Config.from_file(config_path).hyperparam + self.clip = load_clip(name=clip_arch) + + self.config.vpt_layers = list( + range(self.clip.visual.transformer.layers)) + self.config.tpt_layers = list(range(self.clip.transformer.layers)) + + self.pool_frames = BaselinePooling(self.config.pooling_type, + self.config) + + self.visual_prompt_learner = VisualPromptLearner( + self.clip, self.config) + self.image_encoder = ImageEncoder(self.clip, self.config) + + self.text_prompt_learner = TextPromptLearner(self.clip, self.config) + self.text_encoder = TextEncoder(self.clip, self.config) + + # load param from pre-train model + self.load_state_dict(get_state_dict(model_path)) + self.eval() + + # set seed + os.environ['TOKENIZERS_PARALLELISM'] = 'false' + set_seed(self.config.seed) + + def get_video_features(self, videos, return_all_frames=False): + """ + Get video Features + + Args: + videos: the dim is [1, 12, 3, 224, 224] + return_all_frames: default False + """ + batch_size = videos.shape[0] + video_data = videos.reshape(-1, 3, self.config.input_res, + self.config.input_res) + + visual_prompts = self.visual_prompt_learner() + video_features = self.image_encoder(visual_prompts, video_data) + + video_features = video_features / video_features.norm( + dim=-1, keepdim=True) + video_features = video_features.reshape(batch_size, + self.config.num_frames, -1) + + video_features_pooled = self.pool_frames(None, video_features) + + if return_all_frames: + return video_features, video_features_pooled + + return video_features_pooled + + def get_text_features(self, text_data): + """ + Get Text Features + + Args: + text_data: the dim is [1, 69] + """ + text_prompts = self.text_prompt_learner() + text_features = self.text_encoder(text_prompts, text_data) + + text_features = text_features / text_features.norm( + dim=-1, keepdim=True) + return text_features + + def forward(self, data, return_all_frames=False): + """ + Dynamic Forward Function of VoP + + Args: + data: the input data + return_all_frames: default False + """ + batch_size = data['video'].shape[0] + text_data = data['text'] + video_data = data['video'] + video_data = video_data.reshape(-1, 3, self.config.input_res, + self.config.input_res) + + visual_prompts = self.visual_prompt_learner() + video_features = self.image_encoder(visual_prompts, video_data) + + text_prompts = self.text_prompt_learner() + text_features = self.text_encoder(text_prompts, text_data) + + text_features = text_features / text_features.norm( + dim=-1, keepdim=True) + video_features = video_features / video_features.norm( + dim=-1, keepdim=True) + video_features = video_features.reshape(batch_size, + self.config.num_frames, -1) + + video_features_pooled = self.pool_frames(text_features, video_features) + + if return_all_frames: + return text_features, video_features, video_features_pooled + + return text_features, video_features_pooled + + +class BaselinePooling(TorchModel): + """ + Redefined Pooling Function + """ + + def __init__(self, pooling_type, config): + super(BaselinePooling, self).__init__() + if pooling_type == 'avg': + self.pooling_func = self._avg_pooling + else: + raise NotImplementedError + + def _avg_pooling(self, text_embeds, video_embeds): + """ + Pooling mean of frames + + Args: + text_embeds: the input text embedding which is None here. + video_embeds: the input video embedding with [1, 12, 512]. + + Returns: + video_embeds_pooled: num_vids x embed_dim + """ + video_embeds_pooled = video_embeds.mean(dim=1) + return video_embeds_pooled + + def forward(self, text_embeds, video_embeds): + return self.pooling_func(text_embeds, video_embeds) + + +class VisualPromptLearner(TorchModel): + """ + The implementation of visual prompt. + This module is used to define the learnable prompt parameters: + the number of tokens is 8, + the prompt dimension is 768, + and the initialization weight std used is 0.02. + """ + + def __init__(self, clip_model, config): + super(VisualPromptLearner, self).__init__() + + vp_token_num = config.vp_token_num + vp_dim = clip_model.visual.ln_post.weight.shape[0] + dtype = clip_model.dtype + + visual_prompts = torch.empty( + len(config.vpt_layers), 1, vp_token_num, vp_dim, dtype=dtype) + nn.init.normal_(visual_prompts, std=0.02) + self.visual_prompts = nn.Parameter(visual_prompts) + + def forward(self): + vp = self.visual_prompts + return vp + + +class TextPromptLearner(TorchModel): + """ + The implementation of visual prompt. + This module is used to define the learnable prompt parameters: + the number of tokens is 4, + the prompt dimension is 512, + and the initialization weight std used is 0.02. + """ + + def __init__(self, clip_model, config): + super(TextPromptLearner, self).__init__() + + tp_prefix_token_num = config.tp_prefix_token_num + tp_suffix_token_num = config.tp_suffix_token_num + assert tp_prefix_token_num >= 0 and tp_suffix_token_num >= 0 + tp_dim = clip_model.ln_final.weight.shape[0] + dtype = clip_model.dtype + + text_prompts = torch.empty( + len(config.tpt_layers), + tp_prefix_token_num + tp_suffix_token_num, + tp_dim, + dtype=dtype) + nn.init.normal_(text_prompts, std=0.02) + + self.text_prompts = nn.Parameter(text_prompts) + self.tp_prefix_token_num = tp_prefix_token_num + self.tp_suffix_token_num = tp_suffix_token_num + + def forward(self): + return (self.text_prompts[:, :self.tp_prefix_token_num, :], + self.text_prompts[:, self.tp_prefix_token_num:, :]) + + +class ImageEncoder(TorchModel): + """ + The implementation of image encoder. + This module is used to obtain the features of each frame of the video. + """ + + def __init__(self, clip_model, config): + super(ImageEncoder, self).__init__() + + self.config = config + self.vpt_layers = config.vpt_layers + self.vp_token_num = config.vp_token_num + self.num_frames = config.num_frames + + self.conv1 = clip_model.visual.conv1 + self.class_embedding = clip_model.visual.class_embedding + self.positional_embedding = clip_model.visual.positional_embedding + self.ln_pre = clip_model.visual.ln_pre + + self.transformer = clip_model.visual.transformer + + self.ln_post = clip_model.visual.ln_post + self.proj = clip_model.visual.proj + + def forward(self, visual_prompts, x): + """ + The forward function of image encoder. + + Args: + visual_prompts: the visual prompt, dim is [12, 1, 8, 768] + x: the input data, dim is [12, 3, 224, 224] + + Returns: + x: the output data, dim is [12, 512] + """ + batch_size = x.shape[0] + x = self.conv1(x) + x = x.reshape(batch_size, x.shape[1], -1) + x = x.permute(0, 2, 1) + x_1 = self.class_embedding.to(x.dtype) + x_2 = torch.zeros( + batch_size, 1, x.shape[-1], dtype=x.dtype, device=x.device) + x_1 = x_1 + x_2 + x = torch.cat([x_1, x], dim=1) + x = x + self.positional_embedding.to(x.dtype) + + for i_layer in range(self.transformer.layers): + if i_layer in self.vpt_layers: + i_prompt = self.vpt_layers.index(i_layer) + cur_layer_vp = visual_prompts[i_prompt, :, :, :].repeat( + batch_size, 1, 1) + x = torch.cat([x[:, :1, :], cur_layer_vp, x[:, 1:, :]], dim=1) + + if i_layer == 0: + x = self.ln_pre(x) + x = x.permute(1, 0, 2) + x = self.transformer.resblocks[i_layer](x) + x = x.permute(1, 0, 2) + + if i_layer + 1 in self.vpt_layers: + x = torch.cat([x[:, :1, :], x[:, 1 + self.vp_token_num:, :]], + dim=1) + + x = self.ln_post(x[:, 0, :]) + + if self.proj is not None: + x = x @ self.proj + + return x + + +class TextEncoder(TorchModel): + """ + The implementation of text encoder. + This module is used to obtain the features of each word of the sentence. + """ + + def __init__(self, clip_model, config): + super(TextEncoder, self).__init__() + self.transformer = clip_model.transformer + self.token_embedding = clip_model.token_embedding + self.positional_embedding = clip_model.positional_embedding + self.ln_final = clip_model.ln_final + self.text_projection = clip_model.text_projection + self.dtype = clip_model.dtype + + self.tpt_layers = config.tpt_layers + assert 0 in self.tpt_layers + self.tp_prefix_token_num = config.tp_prefix_token_num + self.tp_suffix_token_num = config.tp_suffix_token_num + self.tp_token_num = config.tp_prefix_token_num + config.tp_suffix_token_num + + def forward(self, text_prompts, text): + """ + The forward function of text encoder. + + Args: + text_prompts: the text prompt, dim is 2 x [12, 4, 512] + text: the input data, dim is [1, 69] + + Returns: + x: the output data, dim is [1, 512] + """ + x = self.token_embedding(text).type(self.dtype) + batch_size = x.shape[0] + prompt_prefix, prompt_suffix = text_prompts + + for i_layer in range(self.transformer.layers): + if i_layer in self.tpt_layers: + i_prompt = self.tpt_layers.index(i_layer) + if self.tp_prefix_token_num > 0: + cur_layer_tp_prefix = prompt_prefix[i_prompt:i_prompt + + 1, :, :].expand( + batch_size, -1, -1) + x = torch.cat( + [x[:, :1, :], cur_layer_tp_prefix, x[:, 1:, :]], dim=1) + if self.tp_suffix_token_num > 0: + cur_layer_tp_suffix = prompt_suffix[i_prompt:i_prompt + + 1, :, :].expand( + batch_size, -1, -1) + x = torch.cat( + [x[:, :-1, :], cur_layer_tp_suffix, x[:, -1:, :]], + dim=1) + + if i_layer == 0: + x = x + self.positional_embedding.type(self.dtype) + x = x.permute(1, 0, 2) + x = self.transformer.resblocks[i_layer](x) + x = x.permute(1, 0, 2) + + if i_layer + 1 in self.tpt_layers: + temp_1 = x[:, :1, :] + temp_2 = x[:, 1 + self.tp_prefix_token_num:-1 + - self.tp_suffix_token_num, :] + temp_3 = x[:, -1:, :] + temp = torch.cat([temp_1, temp_2, temp_3], dim=1) + x = temp + + x = self.ln_final(x).type(self.dtype) + x = x[torch.arange(x.shape[0]), + text.argmax(dim=-1) + self.tp_token_num] @ self.text_projection + + return x diff --git a/modelscope/models/cv/vop_retrieval/tokenization_clip.py b/modelscope/models/cv/vop_retrieval/tokenization_clip.py new file mode 100644 index 00000000..07bad10c --- /dev/null +++ b/modelscope/models/cv/vop_retrieval/tokenization_clip.py @@ -0,0 +1,159 @@ +# The implementation here is modified based on HuggingFace, originally Apache 2.0 License +# and publicly avaialbe at https://github.com/huggingface/transformers +# Copyright 2018 The HuggingFace Inc. team. +# Copyright 2021-2022 The Alibaba Fundamental Vision Team Authors. All rights reserved. + +import gzip +import html +import os +from functools import lru_cache + +import ftfy +import regex as re +import torch + + +@lru_cache() +def bytes_to_unicode(): + bs = list(range(ord('!'), + ord('~') + 1)) + list(range( + ord('¡'), + ord('¬') + 1)) + list(range(ord('®'), + ord('ÿ') + 1)) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +def get_pairs(word): + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +class LengthAdaptiveTokenizer(object): + + def __init__(self, config, bpe_path): + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + merges = bpe_path + merges = merges[1:49152 - 256 - 2 + 1] + merges = [tuple(merge.split()) for merge in merges] + vocab = list(bytes_to_unicode().values()) + vocab = vocab + [v + '' for v in vocab] + for merge in merges: + vocab.append(''.join(merge)) + vocab.extend(['<|startoftext|>', '<|endoftext|>']) + self.encoder = dict(zip(vocab, range(len(vocab)))) + self.decoder = {v: k for k, v in self.encoder.items()} + self.bpe_ranks = dict(zip(merges, range(len(merges)))) + self.cache = { + '<|startoftext|>': '<|startoftext|>', + '<|endoftext|>': '<|endoftext|>' + } + self.pat = re.compile( + r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", + re.IGNORECASE) + + self.vocab = self.encoder + + self.tp_token_num = config.tp_prefix_token_num + config.tp_suffix_token_num + + def bpe(self, token): + if token in self.cache: + return self.cache[token] + word = tuple(token[:-1]) + (token[-1] + '', ) + pairs = get_pairs(word) + + if not pairs: + return token + '' + + while True: + bigram = min( + pairs, key=lambda pair: self.bpe_ranks.get(pair, float('inf'))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + new_word.extend(word[i:j]) + i = j + except ValueError: + new_word.extend(word[i:]) + break + + if word[i] == first and i < len(word) - 1 and word[ + i + 1] == second: + new_word.append(first + second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + else: + pairs = get_pairs(word) + word = ' '.join(word) + self.cache[token] = word + return word + + def encode(self, text): + bpe_tokens = [] + text = whitespace_clean(basic_clean(text)).lower() + for token in re.findall(self.pat, text): + token = ''.join(self.byte_encoder[b] + for b in token.encode('utf-8')) + bpe_tokens.extend(self.encoder[bpe_token] + for bpe_token in self.bpe(token).split(' ')) + return bpe_tokens + + def __call__(self, + texts, + return_tensors='pt', + padding=True, + truncation=True): + context_length = 77 - self.tp_token_num + if isinstance(texts, str): + texts = [texts] + + sot_token = self.encoder['<|startoftext|>'] + eot_token = self.encoder['<|endoftext|>'] + all_tokens = [[sot_token] + self.encode(text) + [eot_token] + for text in texts] + result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) + + for i, tokens in enumerate(all_tokens): + if len(tokens) > context_length: + new_tokens = [sot_token + ] + tokens[1:context_length - 1] + [eot_token] + result[i, :len(tokens)] = torch.tensor(new_tokens) + else: + result[i, :len(tokens)] = torch.tensor(tokens) + + return result diff --git a/modelscope/pipelines/cv/__init__.py b/modelscope/pipelines/cv/__init__.py index ae4dd96a..22317700 100644 --- a/modelscope/pipelines/cv/__init__.py +++ b/modelscope/pipelines/cv/__init__.py @@ -72,6 +72,7 @@ if TYPE_CHECKING: from .vision_middleware_pipeline import VisionMiddlewarePipeline from .video_frame_interpolation_pipeline import VideoFrameInterpolationPipeline from .image_skychange_pipeline import ImageSkychangePipeline + from .vop_retrieval_pipeline import VopRetrievalPipeline from .video_object_segmentation_pipeline import VideoObjectSegmentationPipeline from .video_stabilization_pipeline import VideoStabilizationPipeline from .video_super_resolution_pipeline import VideoSuperResolutionPipeline @@ -174,6 +175,7 @@ else: 'VideoFrameInterpolationPipeline' ], 'image_skychange_pipeline': ['ImageSkychangePipeline'], + 'vop_retrieval_pipeline': ['VopRetrievalPipeline'], 'video_object_segmentation_pipeline': [ 'VideoObjectSegmentationPipeline' ], diff --git a/modelscope/pipelines/cv/vop_retrieval_pipeline.py b/modelscope/pipelines/cv/vop_retrieval_pipeline.py new file mode 100644 index 00000000..f0907027 --- /dev/null +++ b/modelscope/pipelines/cv/vop_retrieval_pipeline.py @@ -0,0 +1,122 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import gzip +import math +import os +import os.path as osp +import pickle +import random +from collections import defaultdict, deque +from typing import Any, Dict + +import numpy as np +import torch +from tqdm import tqdm + +from modelscope.metainfo import Pipelines +from modelscope.models import Model +from modelscope.models.cv.vop_retrieval import (LengthAdaptiveTokenizer, VoP, + init_transform_dict, load_data, + load_frames_from_video) +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input, Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.preprocessors import load_image +from modelscope.utils.config import Config +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.vop_retrieval, module_name=Pipelines.vop_retrieval) +class VopRetrievalPipeline(Pipeline): + + def __init__(self, model: str, **kwargs): + """ + use `model` to create a vop pipeline for retrieval + Args: + model: model id on modelscope hub. + """ + super().__init__(model=model, **kwargs) + + # [from pretrain] load model + self.model = Model.from_pretrained('damo/cv_vit-b32_retrieval_vop').to( + self.device) + logger.info('load model done') + + # others: load transform + self.local_pth = model + self.cfg = Config.from_file(osp.join(model, ModelFile.CONFIGURATION)) + self.img_transform = init_transform_dict( + self.cfg.hyperparam.input_res)['clip_test'] + logger.info('load transform done') + + # others: load tokenizer + bpe_path = gzip.open(osp.join( + model, + 'bpe_simple_vocab_16e6.txt.gz')).read().decode('utf-8').split('\n') + self.tokenizer = LengthAdaptiveTokenizer(self.cfg.hyperparam, bpe_path) + logger.info('load tokenizer done') + + # others: load dataset + self.database = load_data( + osp.join(model, 'VoP_msrvtt9k_features.pkl'), self.device) + logger.info('load database done') + + def preprocess(self, input: Input) -> Dict[str, Any]: + if isinstance(input, str): + if '.mp4' in input: + query = [] + for video_path in [input]: + video_path = osp.join(self.local_pth, video_path) + imgs, idxs = load_frames_from_video( + video_path, self.cfg.hyperparam.num_frames, + self.cfg.hyperparam.video_sample_type) + imgs = self.img_transform(imgs) + query.append(imgs) + query = torch.stack( + query, dim=0).to( + self.device, non_blocking=True) + mode = 'v2t' + else: + query = self.tokenizer( + input, return_tensors='pt', padding=True, truncation=True) + if isinstance(query, torch.Tensor): + query = query.to(self.device, non_blocking=True) + else: + query = { + key: val.to(self.device, non_blocking=True) + for key, val in query.items() + } + mode = 't2v' + else: + raise TypeError(f'input should be a str,' + f' but got {type(input)}') + result = {'input_data': query, 'mode': mode} + return result + + def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: + text_embeds, vid_embeds_pooled, vid_ids, texts = self.database + with torch.no_grad(): + if input['mode'] == 't2v': + query_feats = self.model.get_text_features(input['input_data']) + score = query_feats @ vid_embeds_pooled.T + retrieval_idxs = torch.topk( + score, k=self.cfg.hyperparam.topk, + dim=-1)[1].cpu().numpy() + res = np.array(vid_ids)[retrieval_idxs] + elif input['mode'] == 'v2t': + query_feats = self.model.get_video_features( + input['input_data']) + score = query_feats @ text_embeds.T + retrieval_idxs = torch.topk( + score, k=self.cfg.hyperparam.topk, + dim=-1)[1].cpu().numpy() + res = np.array(texts)[retrieval_idxs] + results = {'output_data': res, 'mode': input['mode']} + return results + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index b8afdcb9..992c4453 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -93,6 +93,7 @@ class CVTasks(object): virtual_try_on = 'virtual-try-on' movie_scene_segmentation = 'movie-scene-segmentation' language_guided_video_summarization = 'language-guided-video-summarization' + vop_retrieval = 'video-text-retrieval' # video segmentation video_object_segmentation = 'video-object-segmentation' diff --git a/pose_keypoint.jpg b/pose_keypoint.jpg new file mode 100644 index 00000000..550a892f --- /dev/null +++ b/pose_keypoint.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c18fbde1e9c681ce8927a7b851b366ac71f13a0ffaba1bec202c49e49216d9d3 +size 191543 diff --git a/requirements/cv.txt b/requirements/cv.txt index 7ea00b1e..46f4302e 100644 --- a/requirements/cv.txt +++ b/requirements/cv.txt @@ -38,4 +38,6 @@ tf_slim timm>=0.4.9 torchmetrics>=0.6.2 torchvision +ujson +utils videofeatures_clipit>=1.0 diff --git a/tests/pipelines/test_vop_retrieval.py b/tests/pipelines/test_vop_retrieval.py new file mode 100644 index 00000000..c9c356c5 --- /dev/null +++ b/tests/pipelines/test_vop_retrieval.py @@ -0,0 +1,36 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import unittest + +from modelscope.models import Model +from modelscope.models.cv.vop_retrieval import VoP +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class VopRetrievalTest(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = Tasks.vop_retrieval + # self.model_id = '../cv_vit-b32_retrieval_vop' + self.model_id = 'damo/cv_vit-b32_retrieval_vop' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_modelhub(self): + vop_pipeline = pipeline(self.task, self.model_id) + # t2v + result = vop_pipeline('a squid is talking') + # v2t + # result = vop_pipeline('video10.mp4') + print(f'vop output: {result}.') + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_load_model_from_pretrained(self): + # model = Model.from_pretrained('../cv_vit-b32_retrieval_vop') + model = Model.from_pretrained('damo/cv_vit-b32_retrieval_vop') + self.assertTrue(model.__class__ == VoP) + + +if __name__ == '__main__': + unittest.main()