Merge pull request video-db#25 from video-db/ankit/add-transcript-params

ankit-v2-3 · web-flow · commit f9963ca741fa · 2024-08-08T11:18:28.000+05:30
Ankit/add transcript params
diff --git a/videodb/__about__.py b/videodb/__about__.py
@@ -1,7 +1,7 @@
 """ About information for videodb sdk"""
 
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 __title__ = "videodb"
 __author__ = "videodb"
 __email__ = "contact@videodb.io"
diff --git a/videodb/__init__.py b/videodb/__init__.py
@@ -11,6 +11,7 @@
     SceneExtractionType,
     MediaType,
     SearchType,
+    Segmenter,
     SubtitleAlignment,
     SubtitleBorderStyle,
     SubtitleStyle,
@@ -41,6 +42,7 @@
     "SubtitleStyle",
     "TextStyle",
     "SceneExtractionType",
+    "Segmenter",
 ]
 
 
diff --git a/videodb/_constants.py b/videodb/_constants.py
@@ -36,6 +36,12 @@ class SemanticSearchDefaultValues:
     score_threshold = 0.2
 
 
+class Segmenter:
+    time = "time"
+    word = "word"
+    sentence = "sentence"
+
+
 class ApiPath:
     collection = "collection"
     upload = "upload"
diff --git a/videodb/video.py b/videodb/video.py
@@ -5,6 +5,7 @@
     IndexType,
     SceneExtractionType,
     SearchType,
+    Segmenter,
     SubtitleStyle,
     Workflows,
 )
@@ -124,23 +125,61 @@ def get_thumbnails(self) -> List[Image]:
         )
         return [Image(self._connection, **thumbnail) for thumbnail in thumbnails_data]
 
-    def _fetch_transcript(self, force: bool = False) -> None:
-        if self.transcript and not force:
+    def _fetch_transcript(
+        self,
+        start: int = None,
+        end: int = None,
+        segmenter: str = Segmenter.word,
+        length: int = 1,
+        force: bool = None,
+    ) -> None:
+        if (
+            self.transcript
+            and not start
+            and not end
+            and not segmenter
+            and not length
+            and not force
+        ):
             return
         transcript_data = self._connection.get(
             path=f"{ApiPath.video}/{self.id}/{ApiPath.transcription}",
-            params={"force": "true" if force else "false"},
+            params={
+                "start": start,
+                "end": end,
+                "segmenter": segmenter,
+                "length": length,
+                "force": "true" if force else "false",
+            },
             show_progress=True,
         )
         self.transcript = transcript_data.get("word_timestamps", [])
         self.transcript_text = transcript_data.get("text", "")
 
-    def get_transcript(self, force: bool = False) -> List[Dict]:
-        self._fetch_transcript(force)
+    def get_transcript(
+        self,
+        start: int = None,
+        end: int = None,
+        segmenter: str = Segmenter.word,
+        length: int = 1,
+        force: bool = None,
+    ) -> List[Dict]:
+        self._fetch_transcript(
+            start=start, end=end, segmenter=segmenter, length=length, force=force
+        )
         return self.transcript
 
-    def get_transcript_text(self, force: bool = False) -> str:
-        self._fetch_transcript(force)
+    def get_transcript_text(
+        self,
+        start: int = None,
+        end: int = None,
+        segmenter: str = Segmenter.word,
+        length: int = 1,
+        force: bool = None,
+    ) -> str:
+        self._fetch_transcript(
+            start=start, end=end, segmenter=segmenter, length=length, force=force
+        )
         return self.transcript_text
 
     def index_spoken_words(
