Fix extractor and gating tests

README.md

@@ -219,6 +219,13 @@ video_sampler hash "https://www.youtube.com/watch?v=GbpP3Sxp-1U&list=PLFezMcAw96
   --hash-size 3 --buffer-size 20 --ytdlp --keywords "cat,dog,another keyword,test keyword"
 ```
 
+- segment based on audio transcription
+
+```bash
+video_sampler hash some_video.mp4 ./frames/ \
+  --extractor "audio_keyword" --keywords "alert,bell,voice"
+```
+
 The videos are never directly downloaded, only streamed, so you can use it to sample videos from the internet without downloading them first.
 
 ##### Extra YT-DLP options

tests/test_audio_extractor.py

@@ -0,0 +1,11 @@
+import sys
+
+import pytest
+
+from video_sampler.language.keyword_capture import AudioKeywordExtractor
+
+
+def test_audio_extractor_requires_whisper(monkeypatch):
+    monkeypatch.setitem(sys.modules, "whisper", None)
+    with pytest.raises(ImportError):
+        AudioKeywordExtractor(["test"])

video_sampler/gating.py

@@ -23,6 +23,12 @@
 
 
 def create_model(model_name: str):
+    if "open_clip" not in globals() or "torch" not in globals():
+        raise ImportError(
+            "open_clip and torch are required for the clip gate. "
+            "Install them to use this feature."
+        )
+
     model, _, preprocess = open_clip.create_model_and_transforms(
         model_name, pretrained="laion2b_s34b_b79k"
     )

video_sampler/language/__init__.py

@@ -0,0 +1,3 @@
+from .keyword_capture import AudioKeywordExtractor, KeywordExtractor
+
+__all__ = ["KeywordExtractor", "AudioKeywordExtractor"]

video_sampler/language/keyword_capture.py

@@ -81,18 +81,7 @@ def __init__(self, keywords: list[str]) -> None:
     def generate_segments(
         self, subtitle_list: list[tuple[tuple[int, int], str]]
     ) -> Iterable[subtitle_line]:
-        """
-        Captures keyword segments from a list of subtitles.
-
-        Args:
-            subtitle_list (list[tuple[tuple[int, int], str]]): List of subtitles in the format
-                (start_time, end_time, content).
-
-        Yields:
-            subtitle_line: A named tuple representing a keyword segment in the format
-                (start_time, end_time, lemma, content).
-
-        """
+        """Capture segments containing the configured keywords."""
         for (start_time, end_time), content in subtitle_list:
             doc = self.nlp(content.lower())
             for lemma in doc:
@@ -105,8 +94,37 @@ def generate_segments(
                     break
 
 
+class AudioKeywordExtractor(KeywordExtractor):
+    """Extract keywords from an audio file by transcribing it first."""
+
+    def __init__(self, keywords: list[str], model: str = "base") -> None:
+        super().__init__(keywords)
+        try:
+            import whisper
+        except ImportError as e:  # pragma: no cover - optional dependency
+            raise ImportError(
+                "To use this feature install whisper by 'pip install openai-whisper'"
+            ) from e
+        self.whisper = whisper.load_model(model)
+
+    def transcribe(self, audio_path: str) -> list[tuple[tuple[int, int], str]]:
+        result = self.whisper.transcribe(audio_path)
+        subtitle_list: list[tuple[tuple[int, int], str]] = []
+        for seg in result.get("segments", []):
+            subtitle_list.append(
+                ((int(seg["start"] * 1000), int(seg["end"] * 1000)), seg["text"])
+            )
+        return subtitle_list
+
+    def generate_segments_from_audio(self, audio_path: str) -> Iterable[subtitle_line]:
+        subtitle_list = self.transcribe(audio_path)
+        yield from super().generate_segments(subtitle_list)
+
+
 def create_extractor(config: dict):
     if config["type"] == "keyword":
         return KeywordExtractor(**config["args"])
+    if config["type"] == "audio_keyword":
+        return AudioKeywordExtractor(**config["args"])
 
     raise NotImplementedError(f"{config['type']} not implemented yet")