add vop_se for text video retrival

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11719262

modelscope/metainfo.py

@@ -88,6 +88,7 @@ class Models(object):
     rcp_sceneflow_estimation = 'rcp-sceneflow-estimation'
     image_casmvs_depth_estimation = 'image-casmvs-depth-estimation'
     vop_retrieval_model = 'vop-retrieval-model'
+    vop_retrieval_model_se = 'vop-retrieval-model-se'
     ddcolor = 'ddcolor'
     image_probing_model = 'image-probing-model'
     defrcn = 'defrcn'
@@ -366,6 +367,7 @@ class Pipelines(object):
     image_multi_view_depth_estimation = 'image-multi-view-depth-estimation'
     video_panoptic_segmentation = 'video-panoptic-segmentation'
     vop_retrieval = 'vop-video-text-retrieval'
+    vop_retrieval_se = 'vop-video-text-retrieval-se'
     ddcolor_image_colorization = 'ddcolor-image-colorization'
     image_structured_model_probing = 'image-structured-model-probing'
     image_fewshot_detection = 'image-fewshot-detection'

modelscope/models/cv/vop_retrieval/__init__.py

@@ -6,6 +6,7 @@ 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 .model_se import VoP_SE
     from .tokenization_clip import LengthAdaptiveTokenizer
 else:
     _import_structure = {
@@ -14,6 +15,7 @@ else:
             'load_frames_from_video'
         ],
         'model': ['VoP'],
+        'model_se': ['VideoTextRetrievalModelSeries'],
         'tokenization_clip': ['LengthAdaptiveTokenizer']
     }
 

modelscope/models/cv/vop_retrieval/model_se.py

@@ -0,0 +1,156 @@
+# 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_se)
+class VideoTextRetrievalModelSeries(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).
+                - The pretrain param (ViT-B/32) downloads from OpenAI:
+                - "https://openaipublic.azureedge.net/clip/models/
+                - 40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"
+            - 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(VideoTextRetrievalModelSeries, self).__init__()
+        model_path = osp.join(model_dir, 'VoPSE_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.pool_frames = BaselinePooling(self.config.pooling_type)
+
+        # load param from pre-train model
+        self.load_state_dict(get_state_dict(model_path))
+
+        # eval model
+        self.eval()
+
+    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)
+
+        video_features = self.clip.encode_image(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(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_features = self.clip.encode_text(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)
+
+        text_features = self.clip.encode_text(text_data)
+        video_features = self.clip.encode_image(video_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(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):
+        super(BaselinePooling, self).__init__()
+        if pooling_type == 'avg':
+            self.pooling_func = self._avg_pooling
+        else:
+            raise NotImplementedError
+
+    def _avg_pooling(self, video_embeds):
+        """
+            Pooling mean of frames
+
+            Args:
+                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, video_embeds):
+        return self.pooling_func(video_embeds)

modelscope/pipelines/cv/__init__.py

@@ -84,6 +84,7 @@ if TYPE_CHECKING:
     from .image_skychange_pipeline import ImageSkychangePipeline
     from .image_driving_perception_pipeline import ImageDrivingPerceptionPipeline
     from .vop_retrieval_pipeline import VopRetrievalPipeline
+    from .vop_retrieval_se_pipeline import VopRetrievalSEPipeline
     from .video_object_segmentation_pipeline import VideoObjectSegmentationPipeline
     from .video_deinterlace_pipeline import VideoDeinterlacePipeline
     from .image_matching_pipeline import ImageMatchingPipeline
@@ -225,6 +226,7 @@ else:
             'ImageDrivingPerceptionPipeline'
         ],
         'vop_retrieval_pipeline': ['VopRetrievalPipeline'],
+        'vop_retrieval_se_pipeline': ['VopRetrievalSEPipeline'],
         'video_object_segmentation_pipeline': [
             'VideoObjectSegmentationPipeline'
         ],

modelscope/pipelines/cv/vop_retrieval_se_pipeline.py

@@ -0,0 +1,142 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import gzip
+import os.path as osp
+from typing import Any, Dict
+
+import numpy as np
+import torch
+
+from modelscope.metainfo import Pipelines
+from modelscope.models import Model
+from modelscope.models.cv.vop_retrieval import (LengthAdaptiveTokenizer,
+                                                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.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_se)
+class VopRetrievalSEPipeline(Pipeline):
+
+    def __init__(self, model: str, **kwargs):
+        r""" Card VopRetrievalSE Pipeline.
+
+        Examples:
+        >>>
+        >>>   from modelscope.pipelines import pipeline
+        >>>   vop_pipeline = pipeline(Tasks.vop_retrieval,
+        >>>            model='damo/cv_vit-b32_retrieval_vop_bias')
+        >>>
+        >>>   # IF DO TEXT-TO-VIDEO:
+        >>>   input_text = 'a squid is talking'
+        >>>   result = vop_pipeline(input_text)
+        >>>   result:
+        >>>   {'output_data': array([['video8916']], dtype='<U9'),'mode': 't2v'}
+        >>>
+        >>>   # IF DO VIDEO-TO-TEXT:
+        >>>   input_video = 'video10.mp4'
+        >>>   result = vop_pipeline(input_video)
+        >>>   result:
+        >>>   {'output_data': array([['assorted people are shown holding cute pets']], dtype='<U163'), 'mode': 'v2t'}
+        >>>
+        """
+        super().__init__(model=model, **kwargs)
+
+        # [from pretrain] load model
+        self.model = Model.from_pretrained(model).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
+        if 'vop_bias' in model:
+            self.database = load_data(
+                osp.join(model, 'Bias_msrvtt9k_features.pkl'), self.device)
+        elif 'vop_partial' in model:
+            self.database = load_data(
+                osp.join(model, 'Partial_msrvtt9k_features.pkl'), self.device)
+        elif 'vop_proj' in model:
+            self.database = load_data(
+                osp.join(model, 'Proj_msrvtt9k_features.pkl'), self.device)
+        else:
+            self.database = load_data(
+                osp.join(model, 'VoP_msrvtt9k_features.pkl'), self.device)
+        logger.info('load database done')
+
+    def preprocess(self, input: Input, **preprocess_params) -> 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],
+                **forward_params) -> 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],
+                    **post_params) -> Dict[str, Any]:
+        return inputs

tests/pipelines/test_vop_retrieval_sebias.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 VideoTextRetrievalModelSeries
+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_bias'
+        self.model_id = 'damo/cv_vit-b32_retrieval_vop_bias'
+
+    @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_bias')
+        model = Model.from_pretrained('damo/cv_vit-b32_retrieval_vop_bias')
+        self.assertTrue(model.__class__ == VideoTextRetrievalModelSeries)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/pipelines/test_vop_retrieval_separtial.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 VideoTextRetrievalModelSeries
+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_partial'
+
+    @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_partial')
+        self.assertTrue(model.__class__ == VideoTextRetrievalModelSeries)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/pipelines/test_vop_retrieval_seproj.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 VideoTextRetrievalModelSeries
+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_proj'
+
+    @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_proj')
+        self.assertTrue(model.__class__ == VideoTextRetrievalModelSeries)
+
+
+if __name__ == '__main__':
+    unittest.main()