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.
+
+
+
+
+
+
+## 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()
+```
+
+
+
+
+
+
+
+
+## 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