Merge pull request #12 from video-db/v0.1.0/scene-index

V0.1.0/scene index

Author: ashish-spext

videodb/_constants.py

@@ -14,10 +14,12 @@ class MediaType:
 class SearchType:
     semantic = "semantic"
     keyword = "keyword"
+    scene = "scene"
 
 
 class IndexType:
     semantic = "semantic"
+    scene = "scene"
 
 
 class Workflows:
@@ -44,6 +46,7 @@ class ApiPath:
     compile = "compile"
     workflow = "workflow"
     timeline = "timeline"
+    delete = "delete"
 
 
 class Status:

videodb/_utils/_http_client.py

@@ -133,7 +133,7 @@ def _get_output(self, url: str):
             response_json.get("status") == Status.in_progress
             or response_json.get("status") == Status.processing
         ):
-            percentage = response_json.get("data").get("percentage")
+            percentage = response_json.get("data", {}).get("percentage")
             if percentage and self.show_progress and self.progress_bar:
                 self.progress_bar.n = int(percentage)
                 self.progress_bar.update(0)
@@ -169,7 +169,7 @@ def _parse_response(self, response: requests.Response):
                         bar_format="{l_bar}{bar:100}{r_bar}{bar:-100b}",
                     )
                 response_json = self._get_output(
-                    response_json.get("data").get("output_url")
+                    response_json.get("data", {}).get("output_url")
                 )
                 if response_json.get("success"):
                     return response_json.get("data")

videodb/collection.py

@@ -77,11 +77,11 @@ def search(
     ) -> SearchResult:
         search = SearchFactory(self._connection).get_search(search_type)
         return search.search_inside_collection(
-            self.id,
-            query,
-            result_threshold,
-            score_threshold,
-            dynamic_score_percentage,
+            collection_id=self.id,
+            query=query,
+            result_threshold=result_threshold,
+            score_threshold=score_threshold,
+            dynamic_score_percentage=dynamic_score_percentage,
         )
 
     def upload(

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,39 @@ def search_inside_collection(self, **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,
+        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,
+                "score_threshold": score_threshold,
+                "result_threshold": result_threshold,
+            },
+        )
+        return SearchResult(self._connection, **search_data)
+
+    def search_inside_collection(self, **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,11 @@
-from typing import Optional, List, Dict, Tuple
+from typing import Optional, Union, List, Dict, Tuple
 from videodb._utils._video import play_stream
 from videodb._constants import (
     ApiPath,
-    SearchType,
     IndexType,
-    Workflows,
+    SearchType,
     SubtitleStyle,
+    Workflows,
 )
 from videodb.search import SearchFactory, SearchResult
 from videodb.shot import Shot
@@ -24,6 +24,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 (
@@ -51,11 +52,11 @@ def search(
     ) -> 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,
         )
 
     def delete(self) -> None:
@@ -130,6 +131,43 @@ def index_spoken_words(self) -> None:
             },
         )
 
+    def index_scenes(
+        self,
+        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,
+                "force": force,
+                "prompt": prompt,
+                "callback_url": callback_url,
+            },
+        )
+
+    def get_scenes(self) -> 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,
+            },
+        )
+        self.scenes = scene_data
+        return scene_data if scene_data else None
+
+    def delete_scene_index(self) -> None:
+        self._connection.post(
+            path=f"{ApiPath.video}/{self.id}/{ApiPath.index}/{ApiPath.delete}",
+            data={
+                "index_type": IndexType.scene,
+            },
+        )
+        self.scenes = None
+
     def add_subtitle(self, style: SubtitleStyle = SubtitleStyle()) -> str:
         if not isinstance(style, SubtitleStyle):
             raise ValueError("style must be of type SubtitleStyle")