Add `get_frames_from_dir()`

modules/video_utils.py

@@ -1,18 +1,22 @@
 import subprocess
 import os
+from typing import List, Optional
 
 from modules.logger_util import get_logger
+from modules.paths import TEMP_DIR
 
 logger = get_logger()
 
 
 def extract_frames(
     vid_input: str,
-    output_temp_dir: str,
+    output_temp_dir: str = TEMP_DIR,
     quality: int = 2,
     start_number: int = 0
 ):
-    """Extract frames and save them into output_temp_dir. This needs FFmpeg installed."""
+    """
+    Extract frames and save them into output_temp_dir. This needs FFmpeg installed.
+    """
     os.makedirs(output_temp_dir, exist_ok=True)
     output_path = os.path.join(output_temp_dir, "%05d.jpg")
 
@@ -30,3 +34,18 @@ def extract_frames(
         logger.exception("Error occured while extracting frames from the video")
         raise f"An error occurred: {str(e)}"
 
+
+def get_frames_from_dir(vid_dir: str,
+                        available_extensions: Optional[List, str] = None) -> List:
+    if available_extensions is None:
+        available_extensions = [".jpg", ".jpeg", ".JPG", ".JPEG"]
+
+    if isinstance(available_extensions, str):
+        available_extensions = [available_extensions]
+
+    frame_names = [
+        p for p in os.listdir(vid_dir)
+        if os.path.splitext(p)[-1] in available_extensions
+    ]
+    frame_names.sort(key=lambda x: int(os.path.splitext(x)[0]))
+    return frame_names