feat: add scene index and search

Author: ankit-v2-3

videodb/__init__.py

@@ -8,6 +8,7 @@
 from videodb._constants import (
     VIDEO_DB_API,
     MediaType,
+    SceneModels,
     SearchType,
     SubtitleAlignment,
     SubtitleBorderStyle,
@@ -37,6 +38,7 @@
     "SubtitleAlignment",
     "SubtitleBorderStyle",
     "SubtitleStyle",
+    "SceneModels",
 ]
 
 

videodb/_constants.py

@@ -14,10 +14,18 @@ class MediaType:
 class SearchType:
     semantic = "semantic"
     keyword = "keyword"
+    scene = "scene"
 
 
 class IndexType:
     semantic = "semantic"
+    scene = "scene"
+
+
+class SceneModels:
+    gemini_vision: str = "gemini-vision"
+    gpt4_vision: str = "gpt4-v"
+    all: str = "all"
 
 
 class Workflows:
@@ -44,6 +52,7 @@ class ApiPath:
     compile = "compile"
     workflow = "workflow"
     timeline = "timeline"
+    delete = "delete"
 
 
 class Status:

videodb/search.py

@@ -112,6 +112,7 @@ def search_inside_video(
         result_threshold: Optional[int] = None,
         score_threshold: Optional[int] = None,
         dynamic_score_percentage: Optional[int] = None,
+        **kwargs,
     ):
         search_data = self._connection.post(
             path=f"{ApiPath.video}/{video_id}/{ApiPath.search}",
@@ -133,6 +134,7 @@ def search_inside_collection(
         result_threshold: Optional[int] = None,
         score_threshold: Optional[int] = None,
         dynamic_score_percentage: Optional[int] = None,
+        **kwargs,
     ):
         search_data = self._connection.post(
             path=f"{ApiPath.collection}/{collection_id}/{ApiPath.search}",
@@ -176,7 +178,42 @@ def search_inside_collection(**kwargs):
         raise NotImplementedError("Keyword search will be implemented in the future")
 
 
-search_type = {SearchType.semantic: SemanticSearch, SearchType.keyword: KeywordSearch}
+class SceneSearch(Search):
+    def __init__(self, _connection):
+        self._connection = _connection
+
+    def search_inside_video(
+        self,
+        video_id: str,
+        query: str,
+        scene_model: Optional[str] = None,
+        result_threshold: Optional[int] = None,
+        score_threshold: Optional[int] = None,
+        dynamic_score_percentage: Optional[int] = None,
+        **kwargs,
+    ):
+        search_data = self._connection.post(
+            path=f"{ApiPath.video}/{video_id}/{ApiPath.search}",
+            data={
+                "index_type": SearchType.scene,
+                "query": query,
+                "model_name": scene_model,
+                "score_threshold": score_threshold,
+                "result_threshold": result_threshold,
+            },
+        )
+        print(search_data)
+        return SearchResult(self._connection, **search_data)
+
+    def search_inside_collection(**kwargs):
+        raise NotImplementedError("Scene search will be implemented in the future")
+
+
+search_type = {
+    SearchType.semantic: SemanticSearch,
+    SearchType.keyword: KeywordSearch,
+    SearchType.scene: SceneSearch,
+}
 
 
 class SearchFactory:

videodb/video.py

@@ -1,11 +1,12 @@
-from typing import Optional
+from typing import Optional, Union
 from videodb._utils._video import play_stream
 from videodb._constants import (
     ApiPath,
-    SearchType,
     IndexType,
-    Workflows,
+    SceneModels,
+    SearchType,
     SubtitleStyle,
+    Workflows,
 )
 from videodb.search import SearchFactory, SearchResult
 from videodb.shot import Shot
@@ -24,6 +25,7 @@ def __init__(self, _connection, id: str, collection_id: str, **kwargs) -> None:
         self.length = float(kwargs.get("length", 0.0))
         self.transcript = kwargs.get("transcript", None)
         self.transcript_text = kwargs.get("transcript_text", None)
+        self.scenes = kwargs.get("scenes", None)
 
     def __repr__(self) -> str:
         return (
@@ -45,17 +47,19 @@ def search(
         self,
         query: str,
         search_type: Optional[str] = SearchType.semantic,
+        scene_model: Optional[str] = SceneModels.gemini_vision,
         result_threshold: Optional[int] = None,
         score_threshold: Optional[int] = None,
         dynamic_score_percentage: Optional[int] = None,
     ) -> SearchResult:
         search = SearchFactory(self._connection).get_search(search_type)
         return search.search_inside_video(
-            self.id,
-            query,
-            result_threshold,
-            score_threshold,
-            dynamic_score_percentage,
+            video_id=self.id,
+            query=query,
+            result_threshold=result_threshold,
+            score_threshold=score_threshold,
+            dynamic_score_percentage=dynamic_score_percentage,
+            scene_model=scene_model,
         )
 
     def delete(self) -> None:
@@ -130,6 +134,48 @@ def index_spoken_words(self) -> None:
             },
         )
 
+    def index_scenes(
+        self,
+        scene_model: str = SceneModels.gemini_vision,
+        force: bool = False,
+        prompt: str = None,
+        callback_url: str = None,
+    ) -> None:
+        self._connection.post(
+            path=f"{ApiPath.video}/{self.id}/{ApiPath.index}",
+            data={
+                "index_type": IndexType.scene,
+                "model_name": scene_model,
+                "force": force,
+                "prompt": prompt,
+                "callback_url": callback_url,
+            },
+        )
+
+    def get_scenes(
+        self, scene_model: str = SceneModels.gemini_vision
+    ) -> Union[list, None]:
+        if self.scenes:
+            return self.scenes
+        scene_data = self._connection.get(
+            path=f"{ApiPath.video}/{self.id}/{ApiPath.index}",
+            params={
+                "index_type": IndexType.scene,
+                "model_name": scene_model,
+            },
+        )
+        self.scenes = scene_data
+        return scene_data if scene_data else None
+
+    def delete_scene_index(self, scene_model: str = SceneModels.gemini_vision) -> None:
+        self._connection.post(
+            path=f"{ApiPath.video}/{self.id}/{ApiPath.index}/{ApiPath.delete}",
+            data={
+                "index_type": IndexType.scene,
+                "model_name": scene_model,
+            },
+        )
+
     def add_subtitle(self, style: SubtitleStyle = SubtitleStyle()) -> str:
         if not isinstance(style, SubtitleStyle):
             raise ValueError("style must be of type SubtitleStyle")