forktown/Av1an
1
0
Fork 0
mirror of https://github.com/master-of-zen/Av1an.git synced 2024-11-25 02:29:40 +00:00
Code Issues Projects Releases Packages Wiki Activity

code improvements

Browse source
This commit is contained in:
Zen 2020-08-29 12:56:48 +03:00
parent 65f83d05cc
commit a063558c69
4 changed files with 114 additions and 76 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
Av1an/encode.py
View file

@ -1,7 +1,6 @@
#!/usr/bin/env python3
import time
import json
import os
from pathlib import Path
from typing import List
import sys
@ -10,21 +9,21 @@ import concurrent.futures
import shutil
from Av1an.encoders import ENCODERS
from .arg_parse import Args
from .chunk import Chunk
from .chunk_queue import load_or_gen_chunk_queue
from .concat import concat_routine
from .resume import write_progress_file
from .target_vmaf import target_vmaf_routine
from .utils import frame_probe_cv2, terminate, process_inputs
from .bar import Manager, tqdm_bar
from .setup import determine_resources, outputs_filenames, setup
from .logger import log, set_log
from .config import conf
from .ffmpeg import extract_audio, frame_probe
from .fp_reuse import segment_first_pass
from .split import split_routine
from .vmaf import plot_vmaf
from Av1an.arg_parse import Args
from Av1an.chunk import Chunk
from Av1an.chunk_queue import load_or_gen_chunk_queue
from Av1an.concat import concat_routine
from Av1an.resume import write_progress_file
from Av1an.target_vmaf import target_vmaf_routine
from Av1an.utils import frame_probe_cv2, terminate, process_inputs
from Av1an.bar import Manager, tqdm_bar
from Av1an.setup import determine_resources, outputs_filenames, setup
from Av1an.logger import log, set_log
from Av1an.config import conf
from Av1an.ffmpeg import extract_audio, frame_probe
from Av1an.fp_reuse import segment_first_pass
from Av1an.split import split_routine
from Av1an.vmaf import plot_vmaf
def main_queue(args):
@ -44,6 +43,7 @@ def main_queue(args):
args.output_file = None
encode_file(args)
print(f'Finished: {round(time.time() - tm, 1)}s\n')
except KeyboardInterrupt:
print('Encoding stopped')
@ -58,7 +58,7 @@ def encode_file(args: Args):
:return: None
"""
args.output_file = outputs_filenames(args.input, args.output_file, args.encoder)
outputs_filenames(args)
done_path = args.temp / 'done.json'
resuming = args.resume and done_path.exists()
@ -134,7 +134,6 @@ def encoding_loop(args: Args, chunk_queue: List[Chunk]):
with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor:
future_cmd = {executor.submit(encode, cmd, args): cmd for cmd in chunk_queue}
for future in concurrent.futures.as_completed(future_cmd):
future_cmd[future]
try:
future.result()
except Exception as exc:

50
Av1an/setup.py
View file

@ -11,11 +11,16 @@ from pathlib import Path
from psutil import virtual_memory
from Av1an.encoders import ENCODERS
from .arg_parse import Args
from .utils import terminate
from Av1an.arg_parse import Args
from Av1an.utils import terminate
def set_vmaf(args):
"""
Av1an setup for VMAF
:param args: the Args
"""
if args.vmaf_path:
if not Path(args.vmaf_path).exists():
print(f'No such model: {Path(args.vmaf_path).as_posix()}')
@ -34,13 +39,19 @@ def set_vmaf(args):
def check_exes(args: Args):
"""
Checking required executables
:param args: the Args
"""
if not find_executable('ffmpeg'):
print('No ffmpeg')
terminate()
# this shouldn't happen as encoder choices are validated by argparse
if args.encoder not in ENCODERS:
valid_encoder_str = ", ".join([repr(k) for k in ENCODERS.keys()])
valid_encoder_str = ", ".join([repr(k) for k in ENCODERS])
print(f'Not valid encoder {args.encoder}')
print(f'Valid encoders: {valid_encoder_str}')
terminate()
@ -51,6 +62,11 @@ def check_exes(args: Args):
def setup_encoder(args: Args):
"""
Settup encoder params and passes
:param args: the Args
"""
encoder = ENCODERS[args.encoder]
# validate encoder settings
@ -62,10 +78,15 @@ def setup_encoder(args: Args):
if args.passes is None:
args.passes = encoder.default_passes
args.video_params = encoder.default_args if args.video_params is None else shlex.split(args.video_params)
args.video_params = encoder.default_args if args.video_params is None \
else shlex.split(args.video_params)
def startup_check(args: Args):
"""
Performing essential checks at startup_check
Set constant values
"""
if sys.version_info < (3, 6):
print('Python 3.6+ required')
sys.exit()
@ -80,7 +101,8 @@ def startup_check(args: Args):
set_vmaf(args)
if args.reuse_first_pass and args.encoder != 'aom' and args.split_method != 'aom_keyframes':
print('Reusing the first pass is only supported with the aom encoder and aom_keyframes split method.')
print('Reusing the first pass is only supported with \
the aom encoder and aom_keyframes split method.')
terminate()
setup_encoder(args)
@ -102,7 +124,8 @@ def startup_check(args: Args):
args.ffmpeg = shlex.split(args.ffmpeg)
args.pix_format = ['-strict', '-1', '-pix_fmt', args.pix_format]
args.ffmpeg_pipe = [*args.ffmpeg, *args.pix_format, '-bufsize', '50000K', '-f', 'yuv4mpegpipe', '-']
args.ffmpeg_pipe = [*args.ffmpeg, *args.pix_format,
'-bufsize', '50000K', '-f', 'yuv4mpegpipe', '-']
def determine_resources(encoder, workers):
@ -121,7 +144,7 @@ def determine_resources(encoder, workers):
elif encoder in ('svt_av1', 'svt_vp9', 'x265', 'x264'):
workers = round(min(cpu, ram)) // 8
elif encoder in ('vvc'):
elif encoder in 'vvc':
workers = round(min(cpu, ram))
# fix if workers round up to 0
@ -142,9 +165,12 @@ def setup(temp: Path, resume):
(temp / 'encode').mkdir(exist_ok=True)
def outputs_filenames(inp: Path, out: Path, encoder):
def outputs_filenames(args: Args):
"""
Set output filename
:param args: the Args
"""
suffix = '.mkv'
if out:
return out.with_suffix(suffix)
else:
return Path(f'{inp.stem}_av1{suffix}')
args.output_file = Path(f'{args.output_file}_av1{suffix}') if args.output_file \
else Path(f'{args.input.stem}_av1{suffix}')

51
Av1an/target_vmaf.py
View file

@ -1,24 +1,22 @@
#!/bin/env python
import sys
from math import isnan
import matplotlib
import numpy as np
from matplotlib import pyplot as plt
from scipy import interpolate
from subprocess import STDOUT, PIPE
import subprocess
from subprocess import STDOUT, PIPE
from .arg_parse import Args
from .bar import process_pipe
from .chunk import Chunk
from .commandtypes import CommandPair, Command
from .logger import log
from .utils import terminate
from .vmaf import call_vmaf, read_vmaf_json
from .encoders import ENCODERS
import matplotlib
from matplotlib import pyplot as plt
import numpy as np
from scipy import interpolate
from Av1an.arg_parse import Args
from Av1an.bar import process_pipe
from Av1an.chunk import Chunk
from Av1an.commandtypes import CommandPair, Command
from Av1an.logger import log
from Av1an.vmaf import call_vmaf, read_vmaf_json
def target_vmaf_routine(args: Args, chunk: Chunk):
@ -42,7 +40,8 @@ def gen_probes_names(chunk: Chunk, q):
def probe_cmd(chunk: Chunk, q, ffmpeg_pipe, encoder, vmaf_rate) -> CommandPair:
"""
Generate and return commands for probes at set Q values
These are specifically not the commands that are generated by the user or encoder defaults, since these
These are specifically not the commands that are generated
by the user or encoder defaults, since these
should be faster than the actual encoding commands.
These should not be moved into encoder classes at this point.
"""
@ -52,11 +51,13 @@ def probe_cmd(chunk: Chunk, q, ffmpeg_pipe, encoder, vmaf_rate) -> CommandPair:
probe_name = gen_probes_names(chunk, q).with_suffix('.ivf').as_posix()
if encoder == 'aom':
params = ['aomenc', '--passes=1', '--threads=8', '--end-usage=q', '--cpu-used=6', f'--cq-level={q}']
params = ['aomenc', '--passes=1', '--threads=8',
'--end-usage=q', '--cpu-used=6', f'--cq-level={q}']
cmd = CommandPair(pipe, [*params, '-o', probe_name, '-'])
elif encoder == 'x265':
params = ['x265', '--log-level', '0', '--no-progress', '--y4m', '--preset', 'faster', '--crf', f'{q}']
params = ['x265', '--log-level', '0', '--no-progress',
'--y4m', '--preset', 'faster', '--crf', f'{q}']
cmd = CommandPair(pipe, [*params, '-o', probe_name, '-'])
elif encoder == 'rav1e':
@ -64,26 +65,31 @@ def probe_cmd(chunk: Chunk, q, ffmpeg_pipe, encoder, vmaf_rate) -> CommandPair:
cmd = CommandPair(pipe, [*params, '-o', probe_name, '-'])
elif encoder == 'vpx':
params = ['vpxenc', '--passes=1', '--pass=1', '--codec=vp9', '--threads=8', '--cpu-used=9', '--end-usage=q',
params = ['vpxenc', '--passes=1', '--pass=1', '--codec=vp9',
'--threads=8', '--cpu-used=9', '--end-usage=q',
f'--cq-level={q}']
cmd = CommandPair(pipe, [*params, '-o', probe_name, '-'])
elif encoder == 'svt_av1':
params = ['SvtAv1EncApp', '-i', 'stdin', '--preset', '8', '--rc', '0', '--qp', f'{q}']
params = ['SvtAv1EncApp', '-i', 'stdin',
'--preset', '8', '--rc', '0', '--qp', f'{q}']
cmd = CommandPair(pipe, [*params, '-b', probe_name, '-'])
elif encoder == 'svt_vp9':
params = ['SvtVp9EncApp', '-i', 'stdin', '-enc-mode', '8', '-q', f'{q}']
params = ['SvtVp9EncApp', '-i', 'stdin',
'-enc-mode', '8', '-q', f'{q}']
# TODO: pipe needs to output rawvideo
cmd = CommandPair(pipe, [*params, '-b', probe_name, '-'])
elif encoder == 'x264':
params = ['x264', '--log-level', 'error', '--demuxer', 'y4m', '-', '--no-progress', '--preset', 'slow', '--crf',
params = ['x264', '--log-level', 'error', '--demuxer', 'y4m',
'-', '--no-progress', '--preset', 'slow', '--crf',
f'{q}']
cmd = CommandPair(pipe, [*params, '-o', probe_name, '-'])
return cmd
def make_pipes(ffmpeg_gen_cmd: Command, command: CommandPair):
ffmpeg_gen_pipe = subprocess.Popen(ffmpeg_gen_cmd, stdout=PIPE, stderr=STDOUT)
@ -95,7 +101,6 @@ def make_pipes(ffmpeg_gen_cmd: Command, command: CommandPair):
return pipe
def get_target_q(scores, vmaf_target):
"""
Interpolating scores to get Q closest to target VMAF

54
Av1an/vmaf.py
View file

@ -11,9 +11,11 @@ import numpy as np
from math import log10
from matplotlib import pyplot as plt
from .bar import process_pipe
from .chunk import Chunk
from Av1an.bar import process_pipe
from Av1an.chunk import Chunk
import matplotlib
matplotlib.use('Agg')
def read_vmaf_json(file, percentile):
"""Reads vmaf file with vmaf scores in it and return N percentile score from it.
@ -28,63 +30,65 @@ def read_vmaf_json(file, percentile):
return perc
def call_vmaf(chunk: Chunk, encoded: Path, n_threads, model, res, fl_path: Path = None, vmaf_rate=0):
def call_vmaf(chunk: Chunk, encoded: Path, n_threads, model, res,
fl_path: Path = None, vmaf_rate=0):
cmd = ''
# settings model path
if model:
mod = f":model_path={model}"
else:
mod = ''
mod = f":model_path={model}" if model else ''
# limiting amount of threads for calculation
if n_threads:
n_threads = f':n_threads={n_threads}'
else:
n_threads = ''
n_threads = f':n_threads={n_threads}' if n_threads else ''
if fl_path is None:
fl_path = chunk.fake_input_path.with_name(encoded.stem).with_suffix('.json')
fl = fl_path.as_posix()
# Change framerate of comparison to framerate of probe
if vmaf_rate != 0:
select_frames = f"select=not(mod(n\\,{vmaf_rate})),"
else:
select_frames = ''
select_frames = f"select=not(mod(n\\,{vmaf_rate}))," if vmaf_rate != 0 else ''
# For vmaf calculation both source and encoded segment scaled to 1080
# Also it's required to use -r before both files of vmaf calculation to avoid errors
cmd_in = ('ffmpeg', '-loglevel', 'info', '-y', '-thread_queue_size', '1024', '-hide_banner', '-r', '60', '-i',
encoded.as_posix(), '-r', '60', '-i', '-')
cmd_in = ('ffmpeg', '-loglevel', 'info', '-y', '-thread_queue_size', '1024', '-hide_banner',
'-r', '60', '-i', encoded.as_posix(), '-r', '60', '-i', '-')
filter_complex = ('-filter_complex',)
distorted = f'[0:v]{select_frames}scale={res}:flags=bicubic:force_original_aspect_ratio=decrease,setpts=PTS-STARTPTS[distorted];'
distorted = f'[0:v]{select_frames}scale={res}:flags=bicubic:\
force_original_aspect_ratio=decrease,setpts=PTS-STARTPTS[distorted];'
ref = fr'[1:v]{select_frames}scale={res}:flags=bicubic:force_original_aspect_ratio=decrease,setpts=PTS-STARTPTS[ref];'
ref = fr'[1:v]{select_frames}scale={res}:flags=bicubic:'\
'force_original_aspect_ratio=decrease,setpts=PTS-STARTPTS[ref];'
vmaf_filter = f"[distorted][ref]libvmaf=log_fmt='json':eof_action=endall:log_path={shlex.quote(fl)}{mod}{n_threads}"
vmaf_filter = f"[distorted][ref]libvmaf=log_fmt='json':eof_action=endall:\
log_path={shlex.quote(fl)}{mod}{n_threads}"
cmd_out = ('-f', 'null', '-')
cmd = (*cmd_in, *filter_complex, distorted + ref + vmaf_filter, *cmd_out)
ffmpeg_gen_pipe = subprocess.Popen(chunk.ffmpeg_gen_cmd, stdout=PIPE, stderr=STDOUT)
pipe = subprocess.Popen(cmd, stdin=ffmpeg_gen_pipe.stdout, stdout=PIPE, stderr=STDOUT, universal_newlines=True)
pipe = subprocess.Popen(cmd, stdin=ffmpeg_gen_pipe.stdout,
stdout=PIPE, stderr=STDOUT, universal_newlines=True)
process_pipe(pipe)
return fl_path
def plot_vmaf(source: Path, encoded: Path, args, model, vmaf_res):
"""
Making VMAF plot after encode is done
"""
print('Calculating Vmaf...\r', end='')
fl_path = encoded.with_name(f'{encoded.stem}_vmaflog').with_suffix(".json")
# call_vmaf takes a chunk, so make a chunk of the entire source
ffmpeg_gen_cmd = ['ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-i', source.as_posix(), *args.pix_format,
'-f', 'yuv4mpegpipe', '-']
ffmpeg_gen_cmd = ['ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-i',
source.as_posix(), *args.pix_format, '-f', 'yuv4mpegpipe', '-']
input_chunk = Chunk(args.temp, 0, ffmpeg_gen_cmd, '', 0, 0)
scores = call_vmaf(input_chunk, encoded, 0, model, vmaf_res, fl_path=fl_path)
@ -98,6 +102,10 @@ def plot_vmaf(source: Path, encoded: Path, args, model, vmaf_res):
def plot_vmaf_score_file(scores: Path, plot_path: Path):
"""
Read vmaf json and plot VMAF values for each frame
"""
perc_1 = read_vmaf_json(scores, 1)
perc_25 = read_vmaf_json(scores, 25)
perc_75 = read_vmaf_json(scores, 75)