drl
copied
10 changed files with 226 additions and 1 deletions
@ -1,2 +1,110 @@ |
|||||
# drl |
|
||||
|
# Video-Text Retrieval Embedding with DRL |
||||
|
|
||||
|
*author: Chen Zhang* |
||||
|
|
||||
|
|
||||
|
<br /> |
||||
|
|
||||
|
|
||||
|
|
||||
|
## 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. |
||||
|
|
||||
|
|
||||
|
<br /> |
||||
|
|
||||
|
 |
||||
|
|
||||
|
## 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() |
||||
|
``` |
||||
|
|
||||
|
 |
||||
|
 |
||||
|
|
||||
|
*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() |
||||
|
``` |
||||
|
|
||||
|
 |
||||
|
 |
||||
|
|
||||
|
<br /> |
||||
|
|
||||
|
|
||||
|
|
||||
|
## 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. |
||||
|
|
||||
|
|
||||
|
<br /> |
||||
|
|
||||
|
|
||||
|
|
||||
|
## 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) |
||||
|
|
||||
|
|
||||
|
|
||||
|
|
||||
|
After Width: | Height: | Size: 82 KiB |
@ -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) |
||||
|
|
Binary file not shown.
@ -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 |
@ -0,0 +1,4 @@ |
|||||
|
torchvision |
||||
|
torch |
||||
|
towhee>=0.7.0 |
||||
|
towhee.models>=0.7.0 |
After Width: | Height: | Size: 15 KiB |
After Width: | Height: | Size: 585 KiB |
After Width: | Height: | Size: 18 KiB |
After Width: | Height: | Size: 16 KiB |
Loading…
Reference in new issue