Spaces:
Runtime error
Runtime error
File size: 2,187 Bytes
29f86be f3cdbdc 29f86be cfc3ad2 29f86be cfc3ad2 29f86be cfc3ad2 e0e4152 cfc3ad2 29f86be 9d18995 29f86be 9d18995 29f86be cfc3ad2 f3cdbdc 9d18995 cfc3ad2 6431157 cfc3ad2 9d18995 f3cdbdc e0e4152 f3cdbdc |
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 |
import subprocess
import os
from typing import List, Optional, Union
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,
quality: int = 2,
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,
'-q:v', str(quality),
'-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 f"An error occurred: {str(e)}"
def get_frames_from_dir(vid_dir: str,
available_extensions: Optional[Union[List, str]] = None) -> 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]
return frames
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 f"Error removing {str(e)}"
|