Added easyocr for videoframes

handler.py

@@ -14,6 +14,7 @@ from decord import VideoReader
 from decord import cpu
 
 import timeit
+import easyocr
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -32,6 +33,7 @@ class EndpointHandler:
 
         logging.set_verbosity_debug()
         self.logger = logging.get_logger(__name__)
+        self.reader = easyocr.Reader(["de", "en"])  # Add more languages if needed
         # Check if CUDA (GPU support) is available
         if torch.cuda.is_available():
             self.logger.info("GPU is available for inference.")
@@ -184,6 +186,19 @@ class EndpointHandler:
             self.logger.info(
                 f"Embedding calculation took {embedding_end_time - embedding_start_time} seconds"
             )
+            # Extract text from each frame using EasyOCR
+            self.logger.info("Extracting text from frames.")
+            text_extraction_start_time = timeit.default_timer()
+            frame_texts = [self.reader.readtext(frame, detail=0) for frame in frames]
+            texts_set = set()
+            for text_list in frame_texts:
+                [texts_set.add(text) for text in text_list]
+            video_metadata["extracted_text"] = texts_set
+            text_extraction_end_time = timeit.default_timer()
+            self.logger.info(
+                f"Text extraction took {text_extraction_end_time - text_extraction_start_time} seconds"
+            )
+
             video_metadata["url"] = video_url
             self.logger.info("Returning embeddings and metadata.")
             return frame_embeddings, video_metadata

requirements.txt

@@ -23,4 +23,5 @@ tqdm==4.66.1
 transformers==4.27.2
 typing_extensions==4.8.0
 urllib3==2.0.7
-decord==0.6.0
+decord==0.6.0
+easyocr==1.7.1