diff --git a/README.md b/README.md index 0345c71..8ea4ba0 100755 --- a/README.md +++ b/README.md @@ -94,6 +94,36 @@ Run `python api_example/amr_demo.py` to reproduce the AMR results. **Limitation**: The maximum video duration is **150s** due to the current benchmark datasets. For CPU users, set `feature_name='clip'` because CLIP+Slowfast or CLIP+Slowfast+PANNs features are very slow without GPUs. +## Marengo embeddings (zero-shot, API-based) +Lighthouse also ships an optional, zero-shot moment retrieval predictor backed by +[TwelveLabs Marengo](https://twelvelabs.io) embeddings. It needs **no local checkpoint, feature +files, or GPU**: the video is segmented and embedded server-side, the query is embedded into the +same 512-dim space, and clips are ranked by cosine similarity. This is handy as a training-free +baseline or for videos longer than the 150s benchmark limit. + +Install the extra and set your API key (create one in the +[TwelveLabs dashboard](https://twelvelabs.io)): +``` +pip install 'lighthouse[twelvelabs] @ git+https://github.com/line/lighthouse.git' +export TWELVELABS_API_KEY= +``` +```python +from lighthouse.models import TwelveLabsPredictor + +model = TwelveLabsPredictor(clip_length=2.0) # reads TWELVELABS_API_KEY from the env + +# a local file or a public video URL works +video = model.encode_video('api_example/RoripwjYFp8_60.0_210.0.mp4') + +query = 'A woman wearing a glass is speaking in front of the camera' +prediction = model.predict(query, video) +print(prediction) +""" +{'pred_relevant_windows': [[start, end, score], ...]} # sorted by score, same format as above +""" +``` +Run `python api_example/marengo_demo.py` to reproduce. + ## Gradio demo Run `python gradio_demo/demo.py`. Upload the video and input text query, and click the blue button. For AMR demo, run `python gradio_demo/amr_demo.py`. diff --git a/api_example/marengo_demo.py b/api_example/marengo_demo.py new file mode 100644 index 0000000..9988bcb --- /dev/null +++ b/api_example/marengo_demo.py @@ -0,0 +1,33 @@ +""" +Copyright $today.year LY Corporation + +LY Corporation licenses this file to you 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: + + https://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 os + +from lighthouse.models import TwelveLabsPredictor +from typing import Dict, List, Optional + +# Zero-shot moment retrieval with TwelveLabs Marengo embeddings. +# No local checkpoint or feature files are needed; the video is embedded +# server-side. Set TWELVELABS_API_KEY first (free key: https://twelvelabs.io). +api_key: Optional[str] = os.environ.get('TWELVELABS_API_KEY') +model: TwelveLabsPredictor = TwelveLabsPredictor(api_key=api_key, clip_length=2.0) + +# encode video clips (local file or a public URL) +video: Dict[str, List[List[float]]] = model.encode_video('api_example/RoripwjYFp8_60.0_210.0.mp4') + +# moment retrieval +query: str = 'A woman wearing a glass is speaking in front of the camera' +prediction: Optional[Dict[str, List[List[float]]]] = model.predict(query, video) +print(prediction) diff --git a/lighthouse/marengo_predictor.py b/lighthouse/marengo_predictor.py new file mode 100644 index 0000000..45fdc7e --- /dev/null +++ b/lighthouse/marengo_predictor.py @@ -0,0 +1,174 @@ +import os + +from typing import Dict, List, Optional + +""" +Copyright $today.year LY Corporation + +LY Corporation licenses this file to you 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: + + https://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. +""" + + +class TwelveLabsPredictor: + """ + Zero-shot video moment retrieval backed by TwelveLabs embeddings. + + Unlike the other predictors in :mod:`lighthouse.models`, this one needs no + local checkpoint or feature files. It segments the video on the server side, + embeds each clip and the text query into the same 512-dim embedding space, + and ranks clips by cosine similarity. The output mirrors the rest of the + library: + + {"pred_relevant_windows": [[start, end, score], ...]} + + The embedding model is configurable via the ``model_name`` argument and + defaults to Marengo (``marengo3.0``), so the class can support future + TwelveLabs embedding models without code changes. + + This is an opt-in addition; existing predictors and their behaviour are + unchanged. It depends on the official ``twelvelabs`` SDK, which is only + imported when this class is instantiated, so users who do not use it pay no + import cost. + + A free API key with a generous free tier is available at https://twelvelabs.io. + """ + + DEFAULT_MODEL_NAME: str = 'marengo3.0' + EMBEDDING_DIM: int = 512 + + def __init__( + self, + api_key: Optional[str] = None, + model_name: str = DEFAULT_MODEL_NAME, + clip_length: float = 2.0, + moment_num: int = 10) -> None: + + try: + from twelvelabs import TwelveLabs + except ImportError as e: + raise ImportError( + 'TwelveLabsPredictor requires the twelvelabs SDK. ' + "Install it with: pip install 'lighthouse[twelvelabs]' " + 'or pip install twelvelabs.') from e + + resolved_key = api_key or os.environ.get('TWELVELABS_API_KEY') + if not resolved_key: + raise ValueError( + 'A TwelveLabs API key is required. Pass api_key=... or set the ' + 'TWELVELABS_API_KEY environment variable. Get a free key at ' + 'https://twelvelabs.io.') + + self._client = TwelveLabs(api_key=resolved_key) + self.model_name: str = model_name + self._clip_length: float = clip_length + self._moment_num: int = moment_num + + @staticmethod + def _cosine_similarity( + query: List[float], + clip: List[float]) -> float: + dot = sum(q * c for q, c in zip(query, clip)) + query_norm = sum(q * q for q in query) ** 0.5 + clip_norm = sum(c * c for c in clip) ** 0.5 + if query_norm == 0.0 or clip_norm == 0.0: + return 0.0 + return dot / (query_norm * clip_norm) + + def encode_video( + self, + video_path: Optional[str] = None, + video_url: Optional[str] = None) -> Dict[str, List[List[float]]]: + """ + Embed every clip of a video with TwelveLabs. + + Provide either a local ``video_path`` or a publicly reachable + ``video_url``. Returns a dict with one ``segments`` entry, each being + ``[start_offset_sec, end_offset_sec, *embedding]``, ready to pass to + :meth:`predict`. + """ + if (video_path is None) == (video_url is None): + raise ValueError('Pass exactly one of video_path or video_url.') + + if video_path is not None: + with open(video_path, 'rb') as f: + task = self._client.embed.tasks.create( + model_name=self.model_name, + video_file=f, + video_clip_length=self._clip_length, + video_embedding_scope=['clip']) + else: + task = self._client.embed.tasks.create( + model_name=self.model_name, + video_url=video_url, + video_clip_length=self._clip_length, + video_embedding_scope=['clip']) + + if task.id is None: + raise RuntimeError('The embedding task creation returned no task id.') + task_id: str = task.id + + # wait_for_done lives on the SDK's runtime wrapper class, which mypy + # cannot see through the statically-typed base client. + self._client.embed.tasks.wait_for_done(task_id=task_id) # type: ignore[attr-defined] + result = self._client.embed.tasks.retrieve(task_id=task_id) + + if result.video_embedding is None or result.video_embedding.segments is None: + raise RuntimeError(f'The embedding task {task_id} returned no segments.') + + segments: List[List[float]] = [] + for segment in result.video_embedding.segments: + if segment.float_ is None: + continue + start = float(segment.start_offset_sec or 0.0) + end = float(segment.end_offset_sec or 0.0) + segments.append([start, end] + list(segment.float_)) + return {'segments': segments} + + def _encode_text( + self, + query: str) -> List[float]: + response = self._client.embed.create(model_name=self.model_name, text=query) + if response.text_embedding is None or not response.text_embedding.segments: + raise RuntimeError('TwelveLabs returned no text embedding for the query.') + embedding = response.text_embedding.segments[0].float_ + if embedding is None: + raise RuntimeError('TwelveLabs returned an empty text embedding for the query.') + return list(embedding) + + def predict( + self, + query: str, + inputs: Dict[str, List[List[float]]]) -> Optional[Dict[str, List[List[float]]]]: + """ + Rank the encoded video clips by similarity to ``query``. + + Returns ``{"pred_relevant_windows": [[start, end, score], ...]}`` sorted + by descending score, capped at ``moment_num`` windows, matching the + prediction format of the other Lighthouse predictors. + """ + segments = inputs.get('segments') + if not segments: + print('Error: No encoded clips found. Did you forget to call encode_video()?') + return None + + query_embedding = self._encode_text(query) + + ranked: List[List[float]] = [] + for segment in segments: + start, end = segment[0], segment[1] + clip_embedding = segment[2:] + score = self._cosine_similarity(query_embedding, clip_embedding) + ranked.append([float(f'{start:.4f}'), float(f'{end:.4f}'), float(f'{score:.4f}')]) + + ranked.sort(key=lambda window: window[2], reverse=True) + return {'pred_relevant_windows': ranked[:self._moment_num]} diff --git a/lighthouse/models.py b/lighthouse/models.py index e57fb1a..57914df 100644 --- a/lighthouse/models.py +++ b/lighthouse/models.py @@ -13,6 +13,7 @@ from lighthouse.feature_extractor.vision_encoder import VisionEncoder from lighthouse.feature_extractor.text_encoder import TextEncoder from lighthouse.feature_extractor.audio_encoder import AudioEncoder +from lighthouse.marengo_predictor import TwelveLabsPredictor as TwelveLabsPredictor from typing import Optional, Union, Mapping, Any, Dict, List, Tuple diff --git a/mypy.ini b/mypy.ini index 41e0b02..d5b7b57 100644 --- a/mypy.ini +++ b/mypy.ini @@ -33,3 +33,6 @@ ignore_missing_imports = True [mypy-msclap.*] ignore_missing_imports = True + +[mypy-twelvelabs.*] +ignore_missing_imports = True diff --git a/setup.py b/setup.py index 43b74f6..be1cc26 100644 --- a/setup.py +++ b/setup.py @@ -6,5 +6,6 @@ install_requires=['easydict', 'pandas', 'tqdm', 'pyyaml', 'scikit-learn', 'ffmpeg-python', 'ftfy', 'regex', 'einops', 'fvcore', 'gradio', 'torchlibrosa', 'librosa', 'msclap', 'transformers<=4.51.3', 'numpy<=1.23.5', 'clip@git+https://github.com/openai/CLIP.git'], + extras_require={'twelvelabs': ['twelvelabs>=1.2.8']}, packages=find_packages(exclude=['training']), ) diff --git a/tests/test_marengo.py b/tests/test_marengo.py new file mode 100644 index 0000000..de912ca --- /dev/null +++ b/tests/test_marengo.py @@ -0,0 +1,48 @@ +import os +import pytest + +from lighthouse.models import TwelveLabsPredictor + + +MOMENT_NUM = 10 +EMBEDDING_DIM = 512 + + +def test_predict_ranks_clips_by_cosine_similarity(): + """No-network unit test: predict() must rank clips by query similarity.""" + predictor = TwelveLabsPredictor.__new__(TwelveLabsPredictor) + predictor._moment_num = MOMENT_NUM + + # clip 0 points the same way as the query, clip 1 is orthogonal. + inputs = { + 'segments': [ + [0.0, 2.0, 1.0, 0.0], + [2.0, 4.0, 0.0, 1.0], + ], + } + predictor._encode_text = lambda query: [1.0, 0.0] # type: ignore[method-assign] + + prediction = predictor.predict('a query', inputs) + windows = prediction['pred_relevant_windows'] + + assert len(windows) == 2 + assert windows[0][:2] == [0.0, 2.0], 'the aligned clip should rank first' + assert windows[0][2] == 1.0 + assert windows[1][2] == 0.0 + assert windows[0][2] >= windows[1][2], 'windows must be sorted by descending score' + + +def test_predict_without_segments_returns_none(): + predictor = TwelveLabsPredictor.__new__(TwelveLabsPredictor) + predictor._moment_num = MOMENT_NUM + assert predictor.predict('a query', {'segments': []}) is None + + +@pytest.mark.skipif( + 'TWELVELABS_API_KEY' not in os.environ, + reason='TWELVELABS_API_KEY not set; skipping live Marengo API test.') +def test_marengo_text_embedding_is_512_dim(): + """Live smoke test: a Marengo text embedding is a 512-dim vector.""" + predictor = TwelveLabsPredictor(clip_length=2.0) + embedding = predictor._encode_text('a person walking a dog') + assert len(embedding) == EMBEDDING_DIM