Spaces:
Runtime error
Runtime error
File size: 3,831 Bytes
29f86be f3cdbdc bbbee26 29f86be cfc3ad2 29f86be cfc3ad2 29f86be cfc3ad2 e0e4152 cfc3ad2 29f86be 9d18995 29f86be 9d18995 7ab19f6 29f86be cfc3ad2 bbbee26 9d18995 cfc3ad2 6431157 cfc3ad2 9d18995 bbbee26 9d18995 f3cdbdc 7ab19f6 f3cdbdc e0e4152 f3cdbdc 7ab19f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import subprocess
import os
from typing import List, Optional, Union
from PIL import Image
import numpy as np
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 = TEMP_DIR,
start_number: int = 0
):
"""
Extract frames as jpg files 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")
command = [
'ffmpeg',
'-y', # Enable overwriting
'-i', vid_input,
'-start_number', str(start_number),
f'{output_path}'
]
try:
subprocess.run(command, check=True)
except subprocess.CalledProcessError as e:
logger.exception("Error occurred while extracting frames from the video")
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,
available_extensions: Optional[Union[List, str]] = None,
as_numpy: bool = False) -> List:
"""Get image file paths list from the dir"""
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
]
if not frame_names:
return []
frame_names.sort(key=lambda x: int(os.path.splitext(x)[0]))
frames = [os.path.join(vid_dir, name) for name in frame_names]
if as_numpy:
frames = [np.array(Image.open(frame)) for frame in frames]
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')
for filename in os.listdir(image_dir):
if filename.lower().endswith(image_extensions):
file_path = os.path.join(image_dir, filename)
try:
os.remove(file_path)
except Exception as e:
logger.exception("Error while removing image files")
raise RuntimeError(f"An error occurred: {str(e)}")
|