mirror of
https://github.com/master-of-zen/Av1an.git
synced 2024-11-25 10:40:51 +00:00
269 lines
No EOL
8 KiB
Python
269 lines
No EOL
8 KiB
Python
import json
|
|
import sys
|
|
import os
|
|
import shutil
|
|
from psutil import virtual_memory
|
|
from distutils.spawn import find_executable
|
|
from pathlib import Path
|
|
from av1an.commandtypes import Command
|
|
from av1an.utils import frame_probe_fast, hash_path, terminate
|
|
from av1an.concat import vvc_concat, concatenate_ffmpeg, concatenate_mkvmerge
|
|
from av1an.logger import log
|
|
|
|
class Project(object):
|
|
|
|
def __init__(self, initial_data):
|
|
|
|
# Project info
|
|
self.frames: int = 0
|
|
self.counter = None
|
|
self.is_vs: bool = None
|
|
|
|
# Input/Output/Temp
|
|
self.input: Path = None
|
|
self.temp: Path = None
|
|
self.output_file: Path = None
|
|
self.mkvmerge: bool = None
|
|
self.config = None
|
|
self.webm = None
|
|
|
|
# Splitting
|
|
self.chunk_method: str = None
|
|
self.scenes: Path = None
|
|
self.split_method: str = None
|
|
self.extra_split: int = None
|
|
self.min_scene_len: int = None
|
|
|
|
# PySceneDetect split
|
|
self.threshold: float = None
|
|
|
|
# AOM Keyframe split
|
|
self.reuse_first_pass: bool = None
|
|
|
|
# Encoding
|
|
self.passes = None
|
|
self.video_params: Command = None
|
|
self.encoder: str = None
|
|
self.workers: int = None
|
|
|
|
# FFmpeg params
|
|
self.ffmpeg_pipe: Command = None
|
|
self.ffmpeg: str = None
|
|
self.audio_params = None
|
|
self.pix_format: Command = None
|
|
|
|
# Misc
|
|
self.logging = None
|
|
self.resume: bool = None
|
|
self.no_check: bool = None
|
|
self.keep: bool = None
|
|
self.force: bool = None
|
|
|
|
# Vmaf
|
|
self.vmaf: bool = None
|
|
self.vmaf_path: str = None
|
|
self.vmaf_res: str = None
|
|
|
|
# Target Quality
|
|
self.target_quality: int = None
|
|
self.probes: int = None
|
|
self.min_q: int = None
|
|
self.max_q: int = None
|
|
self.vmaf_plots: bool = None
|
|
self.probing_rate: int = None
|
|
self.n_threads: int = None
|
|
self.vmaf_filter: str = None
|
|
|
|
# VVC
|
|
self.vvc_conf: Path = None
|
|
self.video_dimensions = (None, None)
|
|
self.video_framerate = None
|
|
|
|
# Set all initial values
|
|
self.load_project(initial_data)
|
|
|
|
def load_project(self, initial_data):
|
|
"""
|
|
Loads project attributes to this class
|
|
"""
|
|
# Set all initial values
|
|
for key in initial_data:
|
|
setattr(self, key, initial_data[key])
|
|
|
|
def get_frames(self):
|
|
"""
|
|
Get total frame count of input file, returning total_frames from project if already exists
|
|
"""
|
|
if self.frames > 0:
|
|
return self.frames
|
|
else:
|
|
total = frame_probe_fast(self.input, self.is_vs)
|
|
self.frames = total
|
|
return self.frames
|
|
|
|
def set_frames(self, frame_count: int):
|
|
"""
|
|
Setting total frame count for project
|
|
"""
|
|
self.frames = frame_count
|
|
|
|
def outputs_filenames(self):
|
|
"""
|
|
Set output filename
|
|
|
|
:param project: the Project
|
|
"""
|
|
if self.webm:
|
|
suffix = '.webm'
|
|
else:
|
|
suffix = '.mkv'
|
|
|
|
self.output_file = Path(self.output_file).with_suffix(suffix) if self.output_file \
|
|
else Path(f'{self.input.stem}_{self.encoder}{suffix}')
|
|
|
|
def load_project_from_file(self, path_string):
|
|
"""
|
|
Loads projedt attributes from json to this class
|
|
"""
|
|
pth = Path(path_string)
|
|
with open(pth) as json_data:
|
|
data = json.load(json_data)
|
|
self.load_project(data)
|
|
|
|
def save_project_to_file(self, path_string):
|
|
"""
|
|
Save project attributes from json to this class
|
|
"""
|
|
pth = Path(path_string)
|
|
with open(pth, 'w') as json_data:
|
|
json_data.write(self.save_project())
|
|
|
|
def save_project(self):
|
|
"""
|
|
Returns json of this class, which later can be loaded
|
|
"""
|
|
dt = dict(self.__dict__)
|
|
del dt['input']
|
|
del dt['output_file']
|
|
del dt['temp']
|
|
del dt['vmaf_path']
|
|
del dt['config']
|
|
return json.dumps(dt, indent=4, sort_keys=True)
|
|
|
|
def determine_workers(self):
|
|
"""Returns number of workers that machine can handle with selected encoder."""
|
|
|
|
# If set by user, skip
|
|
if self.workers:
|
|
return self.workers
|
|
|
|
cpu = os.cpu_count()
|
|
ram = round(virtual_memory().total / 2 ** 30)
|
|
|
|
if self.encoder in ('aom', 'rav1e', 'vpx'):
|
|
workers = round(min(cpu / 3, ram / 1.5))
|
|
|
|
elif self.encoder in ('svt_av1', 'svt_vp9', 'x265', 'x264'):
|
|
workers = round(min(cpu, ram)) // 8
|
|
|
|
elif self.encoder in 'vvc':
|
|
workers = round(min(cpu, ram)) // 4
|
|
|
|
# fix if workers round up to 0
|
|
if workers == 0:
|
|
workers = 1
|
|
|
|
self.workers = workers
|
|
|
|
def setup(self):
|
|
"""Creating temporally folders when needed."""
|
|
|
|
if self.temp:
|
|
self.temp = Path(str(self.temp))
|
|
else:
|
|
self.temp = Path('.' + str(hash_path(str(self.input))))
|
|
|
|
# Checking is resume possible
|
|
done_path = self.temp / 'done.json'
|
|
self.resume = self.resume and done_path.exists()
|
|
|
|
if not self.resume:
|
|
if self.temp.is_dir():
|
|
shutil.rmtree(self.temp)
|
|
|
|
(self.temp / 'split').mkdir(parents=True, exist_ok=True)
|
|
(self.temp / 'encode').mkdir(exist_ok=True)
|
|
|
|
def concat_routine(self):
|
|
"""
|
|
Runs the concatenation routine with project
|
|
|
|
:param project: the Project
|
|
:return: None
|
|
"""
|
|
try:
|
|
if self.encoder == 'vvc':
|
|
vvc_concat(self.temp, self.output_file.with_suffix('.h266'))
|
|
elif self.mkvmerge:
|
|
concatenate_mkvmerge(self.temp, self.output_file)
|
|
else:
|
|
concatenate_ffmpeg(self.temp, self.output_file, self.encoder)
|
|
except Exception as e:
|
|
_, _, exc_tb = sys.exc_info()
|
|
print(f'Concatenation failed, error\nAt line: {exc_tb.tb_lineno}\nError:{str(e)}')
|
|
log(f'Concatenation failed, aborting, error: {e}\n')
|
|
terminate()
|
|
|
|
def select_best_chunking_method(self):
|
|
"""
|
|
Selecting best chunking method based on available methods
|
|
"""
|
|
if not find_executable('vspipe'):
|
|
self.chunk_method = 'hybrid'
|
|
log('Set Chunking Method: Hybrid')
|
|
else:
|
|
try:
|
|
import vapoursynth
|
|
plugins = vapoursynth.get_core().get_plugins()
|
|
|
|
if 'com.vapoursynth.ffms2' in plugins:
|
|
log('Set Chunking Method: FFMS2\n')
|
|
self.chunk_method = 'vs_ffms2'
|
|
|
|
elif 'systems.innocent.lsmas' in plugins:
|
|
log('Set Chunking Method: L-SMASH\n')
|
|
self.chunk_method = 'vs_lsmash'
|
|
|
|
except:
|
|
log('Vapoursynth not installed but vspipe reachable\n' +
|
|
'Fallback to Hybrid\n')
|
|
self.chunk_method = 'hybrid'
|
|
|
|
def check_exes(self):
|
|
"""
|
|
Checking required executables
|
|
"""
|
|
|
|
if not find_executable('ffmpeg'):
|
|
print('No ffmpeg')
|
|
terminate()
|
|
|
|
if self.chunk_method in ['vs_ffms2', 'vs_lsmash']:
|
|
if not find_executable('vspipe'):
|
|
print('vspipe executable not found')
|
|
terminate()
|
|
|
|
try:
|
|
import vapoursynth
|
|
plugins = vapoursynth.get_core().get_plugins()
|
|
|
|
if self.chunk_method == 'vs_lsmash' and "systems.innocent.lsmas" not in plugins:
|
|
print('lsmas is not installed')
|
|
terminate()
|
|
|
|
if self.chunk_method == 'vs_ffms2' and "com.vapoursynth.ffms2" not in plugins:
|
|
print('ffms2 is not installed')
|
|
terminate()
|
|
except ModuleNotFoundError:
|
|
print('Vapoursynth is not installed')
|
|
terminate() |