Add video util

modules/utils/video_helper.py

@@ -0,0 +1,272 @@
+import subprocess
+import os
+from typing import List, Optional, Union
+from PIL import Image
+import numpy as np
+from dataclasses import dataclass
+import re
+from pathlib import Path
+
+from modules.utils.constants import SOUND_FILE_EXT, VIDEO_FILE_EXT, IMAGE_FILE_EXT, TRANSPARENT_VIDEO_FILE_EXT
+from modules.utils.paths import TEMP_VIDEO_FRAMES_DIR, TEMP_VIDEO_OUT_FRAMES_DIR
+
+
+@dataclass
+class VideoInfo:
+    num_frames: Optional[int] = None
+    frame_rate: Optional[int] = None
+    duration: Optional[float] = None
+    has_sound: Optional[bool] = None
+    codec: Optional[str] = None
+
+
+def extract_frames(
+    vid_input: str,
+    output_temp_dir: str = TEMP_VIDEO_FRAMES_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,
+        '-qscale:v', '2',
+        '-vf', f'scale=iw:ih',
+        '-start_number', str(start_number),
+        f'{output_path}'
+    ]
+
+    try:
+        subprocess.run(command, check=True)
+    except subprocess.CalledProcessError as e:
+        print("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_VIDEO_FRAMES_DIR,
+):
+    """
+    Extract audio from a video file and save it as a separate sound file. This needs FFmpeg installed.
+    """
+    if Path(vid_input).suffix == ".gif":
+        print("Sound extracting process has passed because gif has no sound")
+        return None
+
+    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:
+        print(f"Warning: Failed to extract sound from the video: {e}")
+
+    return output_path
+
+
+def get_video_info(vid_input: str) -> VideoInfo:
+    """
+    Extract video information using ffmpeg.
+    """
+    command = [
+        'ffmpeg',
+        '-i', vid_input,
+        '-map', '0:v:0',
+        '-c', 'copy',
+        '-f', 'null',
+        '-'
+    ]
+
+    try:
+        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                                encoding='utf-8', errors='replace', check=True)
+        output = result.stderr
+
+        num_frames = None
+        frame_rate = None
+        duration = None
+        has_sound = False
+        codec = None
+
+        for line in output.splitlines():
+            if 'Stream #0:0' in line and 'Video:' in line:
+                fps_match = re.search(r'(\d+(?:\.\d+)?) fps', line)
+                if fps_match:
+                    frame_rate = float(fps_match.group(1))
+
+                codec_match = re.search(r'Video: (\w+)', line)
+                if codec_match:
+                    codec = codec_match.group(1)
+
+            elif 'Duration:' in line:
+                duration_match = re.search(r'Duration: (\d{2}):(\d{2}):(\d{2}\.\d{2})', line)
+                if duration_match:
+                    h, m, s = map(float, duration_match.groups())
+                    duration = h * 3600 + m * 60 + s
+
+            elif 'Stream' in line and 'Audio:' in line:
+                has_sound = True
+
+        if frame_rate and duration:
+            num_frames = int(frame_rate * duration)
+
+        return VideoInfo(
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            duration=duration,
+            has_sound=has_sound,
+            codec=codec
+        )
+
+    except subprocess.CalledProcessError as e:
+        print("Error occurred while getting info from the video")
+        return VideoInfo()
+
+
+def create_video_from_frames(
+    frames_dir: str,
+    frame_rate: Optional[int] = None,
+    sound_path: Optional[str] = None,
+    output_dir: Optional[str] = None,
+    output_mime_type: Optional[str] = None,
+):
+    """
+    Create a video from frames and save it to the output_path. This needs FFmpeg installed.
+    """
+    if not os.path.exists(frames_dir):
+        raise "frames_dir does not exist"
+
+    if output_dir is None:
+        output_dir = TEMP_VIDEO_OUT_FRAMES_DIR
+    os.makedirs(output_dir, exist_ok=True)
+
+    frame_img_mime_type = ".png"
+    pix_format = "yuv420p"
+    vid_codec, audio_codec = "libx264", "aac"
+
+    if output_mime_type is None:
+        output_mime_type = ".mp4"
+
+    output_mime_type = output_mime_type.lower()
+    if output_mime_type == ".mov":
+        pix_format = "yuva444p10le"
+        vid_codec, audio_codec = "prores_ks", "aac"
+
+    elif output_mime_type == ".webm":
+        pix_format = "yuva420p"
+        vid_codec, audio_codec = "libvpx-vp9", "libvorbis"
+
+    elif output_mime_type == ".gif":
+        pix_format = None
+        vid_codec, audio_codec = "gif", None
+
+    num_files = len(os.listdir(output_dir))
+    filename = f"{num_files:05d}{output_mime_type}"
+    output_path = os.path.join(output_dir, filename)
+
+    if sound_path is None:
+        temp_sound = os.path.join(TEMP_VIDEO_FRAMES_DIR, "sound.mp3")
+        if os.path.exists(temp_sound):
+            sound_path = temp_sound
+
+    if frame_rate is None:
+        frame_rate = 25  # Default frame rate for ffmpeg
+
+    command = [
+        'ffmpeg',
+        '-y',
+        '-framerate', str(frame_rate),
+        '-i', os.path.join(frames_dir, f"%05d{frame_img_mime_type}"),
+        '-c:v', vid_codec,
+    ]
+
+    if output_mime_type == ".gif":
+        command += [
+            "-filter_complex", "[0:v] palettegen=reserve_transparent=on [p]; [0:v][p] paletteuse",
+            "-loop", "0"
+        ]
+    else:
+        command += [
+            '-pix_fmt', pix_format
+        ]
+
+    command += [output_path]
+
+    if output_mime_type != ".gif" and sound_path is not None:
+        command += [
+            '-i', sound_path,
+            '-c:a', audio_codec,
+            '-strict', 'experimental',
+            '-b:a', '192k',
+            '-shortest'
+        ]
+    try:
+        subprocess.run(command, check=True)
+    except subprocess.CalledProcessError as e:
+        print("Error occurred while creating video from frames")
+    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: Optional[str] = None):
+    """Removes media files from the video frames directory."""
+    if temp_dir is None:
+        temp_dir = TEMP_VIDEO_FRAMES_DIR
+        temp_out_dir = TEMP_VIDEO_OUT_FRAMES_DIR
+    else:
+        temp_out_dir = os.path.join(temp_dir, "out")
+
+    clean_files_with_extension(temp_dir, SOUND_FILE_EXT)
+    clean_files_with_extension(temp_dir, IMAGE_FILE_EXT)
+    clean_files_with_extension(temp_out_dir, IMAGE_FILE_EXT)
+
+
+def clean_files_with_extension(dir_path: str, extensions: List):
+    """Remove files with the given extensions from the directory."""
+    for filename in os.listdir(dir_path):
+        if filename.lower().endswith(tuple(extensions)):
+            file_path = os.path.join(dir_path, filename)
+            try:
+                os.remove(file_path)
+            except Exception as e:
+                print("Error while removing image files")