Add extract_sound

modules/video_utils.py

@@ -13,7 +13,6 @@ logger = get_logger()
 def extract_frames(
     vid_input: str,
     output_temp_dir: str = TEMP_DIR,
-    quality: int = 2,
     start_number: int = 0
 ):
     """
@@ -26,7 +25,6 @@ def extract_frames(
         'ffmpeg',
         '-y',  # Enable overwriting
         '-i', vid_input,
-        '-q:v', str(quality),
         '-start_number', str(start_number),
         f'{output_path}'
     ]
@@ -35,7 +33,36 @@ def extract_frames(
         subprocess.run(command, check=True)
     except subprocess.CalledProcessError as e:
         logger.exception("Error occurred while extracting frames from the video")
-        raise f"An error occurred: {str(e)}"
+        raise RuntimeError(f"An error occurred: {str(e)}")
+
+    return get_frames_from_dir(output_temp_dir)
+
+
+def extract_sound(
+    vid_input: str,
+    output_temp_dir: str = TEMP_DIR,
+):
+    """
+    Extract audio from a video file and save it as a separate sound file. This needs FFmpeg installed.
+    """
+    os.makedirs(output_temp_dir, exist_ok=True)
+    output_path = os.path.join(output_temp_dir, "sound.mp3")
+
+    command = [
+        'ffmpeg',
+        '-y',  # Enable overwriting
+        '-i', vid_input,
+        '-vn',
+        output_path
+    ]
+
+    try:
+        subprocess.run(command, check=True)
+    except subprocess.CalledProcessError as e:
+        logger.exception("Error occurred while extracting sound from the video")
+        raise RuntimeError(f"An error occurred: {str(e)}")
+
+    return output_path
 
 
 def get_frames_from_dir(vid_dir: str,
@@ -63,6 +90,26 @@ def get_frames_from_dir(vid_dir: str,
     return frames
 
 
+def clean_temp_dir(temp_dir: str = TEMP_DIR):
+    """Removes media files from the directory."""
+    clean_sound_files(temp_dir)
+    clean_image_files(temp_dir)
+
+
+def clean_sound_files(sound_dir: str):
+    """Removes all sound files from the directory."""
+    sound_extensions = ('.mp3', '.wav', '.aac', '.flac', '.ogg', '.m4a', '.wma')
+
+    for filename in os.listdir(sound_dir):
+        if filename.lower().endswith(sound_extensions):
+            file_path = os.path.join(sound_dir, filename)
+            try:
+                os.remove(file_path)
+            except Exception as e:
+                logger.exception("Error while removing sound files")
+                raise RuntimeError(f"Error removing {file_path}: {str(e)}")
+
+
 def clean_image_files(image_dir: str):
     """Removes all image files from the dir"""
     image_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp')
@@ -74,4 +121,4 @@ def clean_image_files(image_dir: str):
                 os.remove(file_path)
             except Exception as e:
                 logger.exception("Error while removing image files")
-                raise f"Error removing {str(e)}"
+                raise RuntimeError(f"An error occurred: {str(e)}")