upagge/roop
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Move temporary files to temp/{target-name}

Browse Source
This commit is contained in:
henryruhs 2023-06-08 17:27:20 +02:00
parent b1c2d414ff
commit 703346da79
2 changed files with 27 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@ -1,3 +1,3 @@
.idea .idea
temp
__pycache__ __pycache__
*.onnx

36
roop/utilities.py
View File

@ -3,12 +3,14 @@ import os
import shutil import shutil
import subprocess import subprocess
from pathlib import Path from pathlib import Path
from typing import List, Any from typing import List
import cv2 import cv2
from PIL import Image
import roop.globals import roop.globals
from PIL import Image
TEMP_FILE = 'temp.mp4'
TEMP_DIRECTORY = 'temp'
def run_ffmpeg(args: List) -> None: def run_ffmpeg(args: List) -> None:
@ -16,7 +18,7 @@ def run_ffmpeg(args: List) -> None:
commands.extend(args) commands.extend(args)
try: try:
subprocess.check_output(commands, stderr=subprocess.STDOUT) subprocess.check_output(commands, stderr=subprocess.STDOUT)
except Exception as exception: except Exception:
pass pass
@ -31,33 +33,41 @@ def detect_fps(source_path: str) -> int:
def extract_frames(target_path: str) -> None: def extract_frames(target_path: str) -> None:
run_ffmpeg(['-i', target_path, get_temp_directory_path(target_path) + os.sep + '%04d.png']) temp_directory_path = get_temp_directory_path(target_path)
run_ffmpeg(['-i', target_path, os.path.join(temp_directory_path, '%04d.png')])
def create_video(target_path: str, fps: int) -> None: def create_video(target_path: str, fps: int) -> None:
run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', roop.globals.video_encoder, '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) temp_directory_path = get_temp_directory_path(target_path)
run_ffmpeg(['-i', os.path.join(temp_directory_path, '%04d.png'), '-framerate', str(fps), '-c:v', roop.globals.video_encoder, '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)])
def restore_audio(target_path: str, output_path: str) -> None: def restore_audio(target_path: str, output_path: str) -> None:
run_ffmpeg(['-i', get_temp_file_path(target_path), '-i', target_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path]) temp_file_path = get_temp_file_path(target_path)
run_ffmpeg(['-i', temp_file_path, '-i', target_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path])
if not os.path.isfile(output_path): if not os.path.isfile(output_path):
move_temp(target_path, output_path) move_temp(target_path, output_path)
def get_temp_frames_paths(target_path: str) -> List: def get_temp_frames_paths(target_path: str) -> List:
return glob.glob(get_temp_directory_path(target_path) + os.sep + '*.png') temp_directory_path = get_temp_directory_path(target_path)
return glob.glob(os.path.join(temp_directory_path, '*.png'))
def get_temp_directory_path(target_path: str) -> str: def get_temp_directory_path(target_path: str) -> str:
return os.path.dirname(target_path) + os.sep + 'temp' filename, _ = os.path.splitext(os.path.basename(target_path))
target_name = os.path.dirname(target_path)
return os.path.join(target_name, TEMP_DIRECTORY, filename)
def get_temp_file_path(target_path: str) -> str: def get_temp_file_path(target_path: str) -> str:
return get_temp_directory_path(target_path) + os.sep + 'temp.mp4' temp_directory_path = get_temp_directory_path(target_path)
return os.path.join(temp_directory_path, TEMP_FILE)
def create_temp(target_path: str) -> None: def create_temp(target_path: str) -> None:
Path(get_temp_directory_path(target_path)).mkdir(exist_ok=True) temp_directory_path = get_temp_directory_path(target_path)
Path(temp_directory_path).mkdir(parents=True, exist_ok=True)
def move_temp(target_path: str, output_path: str) -> None: def move_temp(target_path: str, output_path: str) -> None:
@ -68,8 +78,12 @@ def move_temp(target_path: str, output_path: str) -> None:
def clean_temp(target_path: str) -> None: def clean_temp(target_path: str) -> None:
temp_directory_path = get_temp_directory_path(target_path) temp_directory_path = get_temp_directory_path(target_path)
parent_directory_path = os.path.dirname(temp_directory_path)
parent_directory_name = os.path.basename(parent_directory_path)
if not roop.globals.keep_frames and os.path.isdir(temp_directory_path): if not roop.globals.keep_frames and os.path.isdir(temp_directory_path):
shutil.rmtree(temp_directory_path) shutil.rmtree(temp_directory_path)
if not os.listdir(parent_directory_path) and parent_directory_name == TEMP_DIRECTORY:
os.rmdir(parent_directory_path)
def has_image_extention(image_path: str) -> bool: def has_image_extention(image_path: str) -> bool: