diff --git a/README.md b/README.md index b53dbdd..9f6b3a2 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,110 @@ -# drl +# Video-Text Retrieval Embedding with DRL + +*author: Chen Zhang* + + +
+ + + +## Description + +This operator extracts features for video or text with [DRL(Disentangled Representation Learning for Text-Video Retrieval)](https://arxiv.org/pdf/2203.07111v1.pdf), and then it can get the similarity by Weighted Token-wise Interaction (WTI) module. + + +
+ +![](WTI.png) + +## Code Example + +Load an video from path './demo_video.mp4' to generate a video embedding. + +Read the text 'kids feeding and playing with the horse' to generate a text embedding. + + *Write the pipeline in simplified style*: + +```python +import towhee + +towhee.dc(['./demo_video.mp4']) \ + .video_decode.ffmpeg(sample_type='uniform_temporal_subsample', args={'num_samples': 12}) \ + .runas_op(func=lambda x: [y for y in x]) \ + .drl(base_encoder='clip_vit_b32', modality='video', device='cpu') \ + .show() + +towhee.dc(['kids feeding and playing with the horse']) \ + .drl(base_encoder='clip_vit_b32', modality='text', device='cpu') \ + .show() +``` + +![](vect_simplified_video.png) +![](vect_simplified_text.png) + +*Write a same pipeline with explicit inputs/outputs name specifications:* + +```python +import towhee + +towhee.dc['path'](['./demo_video.mp4']) \ + .video_decode.ffmpeg['path', 'frames'](sample_type='uniform_temporal_subsample', args={'num_samples': 12}) \ + .runas_op['frames', 'frames'](func=lambda x: [y for y in x]) \ + .drl['frames', 'vec'](base_encoder='clip_vit_b32', modality='video', device='cpu') \ + .show(formatter={'path': 'video_path'}) + +towhee.dc['text'](['kids feeding and playing with the horse']) \ + .drl['text','vec'](base_encoder='clip_vit_b32', modality='text', device='cpu') \ + .select['text', 'vec']() \ + .show() +``` + +![](vect_explicit_video.png) +![](vect_explicit_text.png) + +
+ + + +## Factory Constructor + +Create the operator via the following factory method + +***drl(base_encoder, modality)*** + +**Parameters:** + +​ ***base_encoder:*** *str* + +​ The base CLIP encode name in DRL model. Supported model names: +- clip_vit_b32 + + +​ ***modality:*** *str* + +​ Which modality(*video* or *text*) is used to generate the embedding. + + +
+ + + +## Interface + +An video-text embedding operator takes a list of [towhee VideoFrame](link/to/towhee/image/api/doc) or string as input and generate an embedding in ndarray. + + +**Parameters:** + +​ ***data:*** *List[towhee.types.VideoFrame]* or *str* + +​ The data (list of VideoFrame(which is uniform subsampled from a video) or text based on specified modality) to generate embedding. + + + +**Returns:** *numpy.ndarray* + +​ The data embedding extracted by model. When text, the shape is (text_token_num, model_dim), when video, the shape is (video_token_num, model_dim) + + + diff --git a/WTI.png b/WTI.png new file mode 100644 index 0000000..d834ef2 Binary files /dev/null and b/WTI.png differ diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..c328b58 --- /dev/null +++ b/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2021 Zilliz. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .drl import DRL + + +def drl(base_encoder: str, modality: str, **kwargs): + return DRL(base_encoder, modality, **kwargs) + diff --git a/demo_video.mp4 b/demo_video.mp4 new file mode 100755 index 0000000..e6fb645 Binary files /dev/null and b/demo_video.mp4 differ diff --git a/drl.py b/drl.py new file mode 100644 index 0000000..96e6727 --- /dev/null +++ b/drl.py @@ -0,0 +1,93 @@ +# Copyright 2021 Zilliz. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch + +from typing import List, Union +from torchvision import transforms +from towhee.models.clip4clip import convert_tokens_to_id +from towhee.operator.base import NNOperator +from towhee import register +from towhee.models import drl, clip4clip +from PIL import Image as PILImage +from towhee.types import VideoFrame +from pathlib import Path + + +@register(output_schema=['vec']) +class DRL(NNOperator): + """ + DRL multi-modal embedding operator + """ + + def __init__(self, base_encoder: str, modality: str, weight_path: str = None, device: str = None): + super().__init__() + self.modality = modality + if weight_path is None: + weight_path = str(Path(__file__).parent / 'clip_vit_b32_wti.pth') + if device is None: + self.device = "cuda" if torch.cuda.is_available() else "cpu" + else: + self.device = device + self.model = drl.create_model(base_encoder=base_encoder, pretrained=True, cdcr=0, weights_path=weight_path) + + self.tokenize = clip4clip.SimpleTokenizer() + self.tfms = transforms.Compose([ + transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize( + (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)) + ]) + self.model.eval() + + def __call__(self, data: Union[str, List[VideoFrame]]): + if self.modality == 'video': + vec = self._inference_from_video(data) + elif self.modality == 'text': + vec = self._inference_from_text(data) + else: + raise ValueError("modality[{}] not implemented.".format(self._modality)) + return vec + + def _inference_from_text(self, text: str): + self.model.eval() + text_ids = convert_tokens_to_id(self.tokenize, text) + text_ids = torch.tensor(text_ids).unsqueeze(0).to(self.device) + text_features = self.model.get_text_feat(text_ids) # B(1), N_t, D + return text_features.squeeze(0).detach().cpu().numpy() # N_t, D + + def _inference_from_video(self, img_list: List[VideoFrame]): + self.model.eval() + max_frames = 12 + video = np.zeros((1, max_frames, 1, 3, 224, 224), dtype=np.float64) + slice_len = len(img_list) + max_video_length = 0 if 0 > slice_len else slice_len + for i, img in enumerate(img_list): + pil_img = PILImage.fromarray(img, img.mode) + tfmed_img = self.tfms(pil_img).unsqueeze(0) + if slice_len >= 1: + video[0, i, ...] = tfmed_img.cpu().numpy() + video_mask = np.zeros((1, max_frames), dtype=np.int32) + video_mask[0, :max_video_length] = [1] * max_video_length + + video = torch.as_tensor(video).float().to(self.device) + pair, bs, ts, channel, h, w = video.shape + video = video.view(pair * bs * ts, channel, h, w) + video_mask = torch.as_tensor(video_mask).float().to(self.device) + + visual_output = self.model.get_video_feat(video, video_mask, shaped=True) # B(1), N_v, D + + return visual_output.squeeze(0).detach().cpu().numpy() # N_v, D diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..3b2f37c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +torchvision +torch +towhee>=0.7.0 +towhee.models>=0.7.0 \ No newline at end of file diff --git a/vect_explicit_text.png b/vect_explicit_text.png new file mode 100644 index 0000000..521ea68 Binary files /dev/null and b/vect_explicit_text.png differ diff --git a/vect_explicit_video.png b/vect_explicit_video.png new file mode 100644 index 0000000..5aabbba Binary files /dev/null and b/vect_explicit_video.png differ diff --git a/vect_simplified_text.png b/vect_simplified_text.png new file mode 100644 index 0000000..5e2eacb Binary files /dev/null and b/vect_simplified_text.png differ diff --git a/vect_simplified_video.png b/vect_simplified_video.png new file mode 100644 index 0000000..ec824f4 Binary files /dev/null and b/vect_simplified_video.png differ