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

Merge pull request #26 from nicomem/using-paths

Browse source 
Use pathlib Paths + type annotations
This commit is contained in:
Zen 2020-02-11 00:54:28 +02:00 committed by GitHub
commit 6c5caef5fa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 69 additions and 72 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -50,7 +50,7 @@ With your own parameters:
4 rav1e workers is optimal for 6/12 cpu 4 rav1e workers is optimal for 6/12 cpu
-s --scenes Path to file with scenes timestamps. -s --scenes Path to file with scenes timestamps.
If given `0` spliting will be ingnored If given `0` spliting will be ignored
If file not exist, new will be generated in current folder If file not exist, new will be generated in current folder
First run to generate stamps, all next reuse it. First run to generate stamps, all next reuse it.
Example: `-s scenes` Example: `-s scenes`

139
av1an.py
View file

@ -11,11 +11,12 @@ import shutil
from os.path import join from os.path import join
from psutil import virtual_memory from psutil import virtual_memory
import argparse import argparse
from shutil import rmtree
from math import ceil from math import ceil
from multiprocessing import Pool from multiprocessing import Pool
import multiprocessing import multiprocessing
import subprocess import subprocess
from pathlib import Path
from typing import Optional, Union
try: try:
import scenedetect import scenedetect
@ -37,7 +38,7 @@ if sys.version_info < (3, 7):
class Av1an: class Av1an:
def __init__(self): def __init__(self):
self.here = os.getcwd() self.temp_dir = Path('.temp')
self.FFMPEG = 'ffmpeg -hide_banner -loglevel error' self.FFMPEG = 'ffmpeg -hide_banner -loglevel error'
self.pix_format = 'yuv420p' self.pix_format = 'yuv420p'
self.encoder = 'aom' self.encoder = 'aom'
@ -50,9 +51,10 @@ class Av1an:
self.logging = None self.logging = None
self.args = None self.args = None
self.encoding_params = '' self.encoding_params = ''
self.output_file = '' self.output_file: Optional[Path] = None
self.pyscene = '' self.pyscene = ''
self.scenes = '' self.scenes: Optional[Path] = None
self.skip_scenes = False
def call_cmd(self, cmd, capture_output=False): def call_cmd(self, cmd, capture_output=False):
if capture_output: if capture_output:
@ -68,14 +70,14 @@ class Av1an:
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--mode', '-m', type=int, default=self.mode, help='Mode 0 - video, Mode 1 - image') parser.add_argument('--mode', '-m', type=int, default=self.mode, help='Mode 0 - video, Mode 1 - image')
parser.add_argument('--encoding_params', '-e', type=str, default=self.encoding_params, help='encoding settings') parser.add_argument('--encoding_params', '-e', type=str, default=self.encoding_params, help='encoding settings')
parser.add_argument('--file_path', '-i', type=str, default='bruh.mp4', help='Input File', required=True) parser.add_argument('--file_path', '-i', type=Path, default='bruh.mp4', help='Input File', required=True)
parser.add_argument('--encoder', '-enc', type=str, default=self.encoder, help='Choosing encoder') parser.add_argument('--encoder', '-enc', type=str, default=self.encoder, help='Choosing encoder')
parser.add_argument('--workers', '-t', type=int, default=0, help='Number of workers') parser.add_argument('--workers', '-t', type=int, default=0, help='Number of workers')
parser.add_argument('--audio_params', '-a', type=str, default='-c:a copy', help='FFmpeg audio settings') parser.add_argument('--audio_params', '-a', type=str, default='-c:a copy', help='FFmpeg audio settings')
parser.add_argument('--threshold', '-tr', type=float, default=self.threshold, help='PySceneDetect Threshold') parser.add_argument('--threshold', '-tr', type=float, default=self.threshold, help='PySceneDetect Threshold')
parser.add_argument('--logging', '-log', type=str, default=self.logging, help='Enable logging') parser.add_argument('--logging', '-log', type=str, default=self.logging, help='Enable logging')
parser.add_argument('--encode_pass', '-p', type=int, default=self.encode_pass, help='Specify encoding passes') parser.add_argument('--encode_pass', '-p', type=int, default=self.encode_pass, help='Specify encoding passes')
parser.add_argument('--output_file', '-o', type=str, default='', help='Specify output file') parser.add_argument('--output_file', '-o', type=Path, default=None, help='Specify output file')
parser.add_argument('--ffmpeg_com', '-ff', type=str, default='', help='FFmpeg commands') parser.add_argument('--ffmpeg_com', '-ff', type=str, default='', help='FFmpeg commands')
parser.add_argument('--pix_format', '-fmt', type=str, default=self.pix_format, help='FFmpeg pixel format') parser.add_argument('--pix_format', '-fmt', type=str, default=self.pix_format, help='FFmpeg pixel format')
parser.add_argument('--scenes', '-s', type=str, default=self.scenes, help='File location for scenes') parser.add_argument('--scenes', '-s', type=str, default=self.scenes, help='File location for scenes')
@ -84,7 +86,12 @@ class Av1an:
self.args = parser.parse_args() self.args = parser.parse_args()
# Set scenes if provided # Set scenes if provided
self.scenes = self.args.scenes if self.args.scenes:
scenes = self.args.scenes.strip()
if scenes == '0':
self.skip_scenes = True
else:
self.scenes = Path(scenes)
self.threshold = self.args.threshold self.threshold = self.args.threshold
@ -100,6 +107,12 @@ class Av1an:
# Number of encoder passes # Number of encoder passes
self.encode_pass = self.args.encode_pass self.encode_pass = self.args.encode_pass
# Set output file
if self.args.output_file is None:
self.output_file = Path(f'{self.args.file_path.stem}_av1.mkv')
else:
self.output_file = self.args.output_file.with_suffix('.mkv')
# Forcing FPS option # Forcing FPS option
if self.args.ffmpeg_com == 0: if self.args.ffmpeg_com == 0:
self.ffmpeg_com = '' self.ffmpeg_com = ''
@ -140,48 +153,49 @@ class Av1an:
if self.workers == 0: if self.workers == 0:
self.workers += 1 self.workers += 1
def setup(self, input_file): def setup(self, input_file: Path):
if not os.path.exists(input_file): if not input_file.exists():
print("File don't exist") print("File don't exist")
sys.exit() sys.exit()
# Make temporal directories, and remove them if already presented # Make temporal directories, and remove them if already presented
if os.path.isdir(join(os.getcwd(), ".temp")): if self.temp_dir.is_dir():
rmtree(join(self.here, ".temp")) shutil.rmtree(self.temp_dir)
os.makedirs(join(self.here, '.temp', 'split')) (self.temp_dir / 'split').mkdir(parents=True)
os.makedirs(join(self.here, '.temp', 'encode')) (self.temp_dir / 'encode').mkdir()
def extract_audio(self, input_vid): def extract_audio(self, input_vid: Path):
# Extracting audio from video file # Extracting audio from video file
# Encoding audio if needed # Encoding audio if needed
ffprobe = 'ffprobe -hide_banner -loglevel error -show_streams -select_streams a' ffprobe = 'ffprobe -hide_banner -loglevel error -show_streams -select_streams a'
# Capture output to check if audio is present # Capture output to check if audio is present
check = fr'{ffprobe} -i {join(self.here,input_vid)}' check = fr'{ffprobe} -i {input_vid}'
is_audio_here = len(self.call_cmd(check, capture_output=True)) > 0 is_audio_here = len(self.call_cmd(check, capture_output=True)) > 0
if is_audio_here: if is_audio_here:
cmd = f'{self.FFMPEG} -i {join(self.here, input_vid)} -vn ' \ cmd = f'{self.FFMPEG} -i {input_vid} -vn ' \
f'{self.args.audio_params} {join(os.getcwd(), ".temp", "audio.mkv")}' f'{self.args.audio_params} {self.temp_dir / "audio.mkv"}'
self.call_cmd(cmd) self.call_cmd(cmd)
def scenedetect(self, video): def scenedetect(self, video: Path):
# PySceneDetect used split video by scenes and pass it to encoder # Skip scene detection if the user choosed to
# Optimal threshold settings 15-50 if self.skip_scenes:
video_manager = VideoManager([video]) return ''
scene_manager = SceneManager()
scene_manager.add_detector(ContentDetector(threshold=self.threshold))
base_timecode = video_manager.get_base_timecode()
try: try:
if self.scenes.isdigit(): # PySceneDetect used split video by scenes and pass it to encoder
if int(self.scenes) == 0: # Optimal threshold settings 15-50
return '' video_manager = VideoManager([str(video)])
scene_manager = SceneManager()
scene_manager.add_detector(ContentDetector(threshold=self.threshold))
base_timecode = video_manager.get_base_timecode()
# If stats file exists, load it. # If stats file exists, load it.
if self.scenes and os.path.exists(join(self.here, self.scenes)): if self.scenes and self.scenes.exists():
# Read stats from CSV file opened in read mode: # Read stats from CSV file opened in read mode:
with open(join(self.here, self.scenes), 'r') as stats_file: with open(self.scenes(), 'r') as stats_file:
stats = stats_file.read() stats = stats_file.read()
return stats return stats
@ -210,11 +224,9 @@ class Av1an:
# We only write to the stats file if a save is required: # We only write to the stats file if a save is required:
if self.scenes: if self.scenes:
with open(join(self.here, self.scenes), 'w') as stats_file: with open(self.scenes, 'w') as stats_file:
stats_file.write(scenes) stats_file.write(scenes)
return scenes return scenes
else:
return scenes
except Exception: except Exception:
print('Error in PySceneDetect') print('Error in PySceneDetect')
sys.exit() sys.exit()
@ -225,25 +237,17 @@ class Av1an:
# If video is single scene, just copy video # If video is single scene, just copy video
if len(timecodes) == 0: if len(timecodes) == 0:
cmd = f'{self.FFMPEG} -i {video} -map_metadata 0 -an -c copy -avoid_negative_ts 1 .temp/split/0.mkv' cmd = f'{self.FFMPEG} -i {video} -map_metadata 0 -an -c copy -avoid_negative_ts 1 {self.temp_dir / "split" / "0.mkv"}'
else: else:
cmd = f'{self.FFMPEG} -i {video} -map_metadata 0 -an -f segment -segment_times {timecodes} ' \ cmd = f'{self.FFMPEG} -i {video} -map_metadata 0 -an -f segment -segment_times {timecodes} ' \
f'-c copy -avoid_negative_ts 1 .temp/split/%04d.mkv' f'-c copy -avoid_negative_ts 1 {self.temp_dir / "split" / "%04d.mkv"}'
self.call_cmd(cmd) self.call_cmd(cmd)
def get_video_queue(self, source_path): def get_video_queue(self, source_path: Path):
# Returns sorted list of all videos that need to be encoded. Big first # Returns sorted list of all videos that need to be encoded. Big first
return sorted(source_path.iterdir(), key=lambda f: -f.stat().st_size)
videos = []
for root, dirs, files in os.walk(source_path):
for file in files:
f = os.path.getsize(os.path.join(root, file))
videos.append([file, f])
videos = [i[0] for i in sorted(videos, key=lambda x: -x[1])]
return videos
def svt_av1_encode(self, file_paths): def svt_av1_encode(self, file_paths):
@ -293,7 +297,7 @@ class Av1an:
if self.encode_pass == 2: if self.encode_pass == 2:
pass_2_commands = [ pass_2_commands = [
(f'-i {file[0]} {self.ffmpeg_pipe}' + (f'-i {file[0]} {self.ffmpeg_pipe}' +
f' {two_pass_1_aom} {self.encoding_params} --fpf={file[0]}.log -o /dev/null - ', f' {two_pass_1_aom} {self.encoding_params} --fpf={file[0]}.log -o {os.devnull} - ',
f'-i {file[0]} {self.ffmpeg_pipe}' + f'-i {file[0]} {self.ffmpeg_pipe}' +
f' {two_pass_2_aom} {self.encoding_params} --fpf={file[0]}.log -o {file[1]}.ivf - ', f' {two_pass_2_aom} {self.encoding_params} --fpf={file[0]}.log -o {file[1]}.ivf - ',
file[2]) file[2])
@ -329,10 +333,9 @@ class Av1an:
return pass_2_commands return pass_2_commands
def compose_encoding_queue(self, files): def compose_encoding_queue(self, files):
file_paths = [(f'{self.temp_dir / "split" / file.name}',
file_paths = [(f'{join(os.getcwd(), ".temp", "split", file_name)}', f'{self.temp_dir / "encode" / file.name}',
f'{join(os.getcwd(), ".temp", "encode", file_name)}', str(file)) for file in files]
file_name) for file_name in files]
if self.encoder == 'aom': if self.encoder == 'aom':
return self.aom_encode(file_paths) return self.aom_encode(file_paths)
@ -357,31 +360,25 @@ class Av1an:
cmd = rf'{self.FFMPEG} {i}' cmd = rf'{self.FFMPEG} {i}'
self.call_cmd(cmd) self.call_cmd(cmd)
def concatenate_video(self, input_video): def concatenate_video(self):
# Using FFMPEG to concatenate all encoded videos to 1 file. # Using FFMPEG to concatenate all encoded videos to 1 file.
# Reading all files in A-Z order and saving it to concat.txt # Reading all files in A-Z order and saving it to concat.txt
with open(f'{join(self.here, ".temp", "concat.txt")}', 'w') as f: concat = self.temp_dir / "concat.txt"
with open(f'{concat}', 'w') as f:
for root, firs, files in os.walk(join(self.here, '.temp', 'encode')): # Write all files that need to be concatenated
for file in sorted(files): # Their path must be relative to the directory where "concat.txt" is
f.write(f"file '{join(root, file)}'\n") encode_files = sorted((self.temp_dir / 'encode').iterdir())
f.writelines(f"file '{file.relative_to(self.temp_dir)}'\n" for file in encode_files)
concat = join(self.here, ".temp", "concat.txt")
# Add the audio file if one was extracted from the input # Add the audio file if one was extracted from the input
audio_file = join(self.here, ".temp", "audio.mkv") audio_file = self.temp_dir / "audio.mkv"
if os.path.exists(audio_file): if audio_file.exists():
audio = f'-i {audio_file} -c:a copy' audio = f'-i {audio_file} -c:a copy'
else: else:
audio = '' audio = ''
if self.output_file == self.args.output_file:
self.output_file = f'{input_video.split(".")[0]}_av1.mkv'
else:
self.output_file = f'{join(self.here, self.args.output_file)}.mkv'
try: try:
cmd = f'{self.FFMPEG} -f concat -safe 0 -i {concat} {audio} -c copy -y {self.output_file}' cmd = f'{self.FFMPEG} -f concat -safe 0 -i {concat} {audio} -c copy -y {self.output_file}'
self.call_cmd(cmd) self.call_cmd(cmd)
@ -390,11 +387,11 @@ class Av1an:
print('Concatenation failed') print('Concatenation failed')
sys.exit() sys.exit()
def image(self, image_path): def image(self, image_path: Path):
print('Encoding Image..', end='') print('Encoding Image..', end='')
image_pipe = rf'{self.FFMPEG} -i {image_path} -pix_fmt yuv420p10le -f yuv4mpegpipe -strict -1 - | ' image_pipe = rf'{self.FFMPEG} -i {image_path} -pix_fmt yuv420p10le -f yuv4mpegpipe -strict -1 - | '
output = rf'{"".join(image_path.split(".")[:-1])}.ivf' output = image_path.with_suffix('.ivf')
if self.encoder == 'aom': if self.encoder == 'aom':
aom = ' aomenc --passes=1 --pass=1 --end-usage=q -b 10 --input-bit-depth=10 ' aom = ' aomenc --passes=1 --pass=1 --end-usage=q -b 10 --input-bit-depth=10 '
cmd = (rf' {image_pipe} ' + cmd = (rf' {image_pipe} ' +
@ -421,7 +418,7 @@ class Av1an:
# Splitting video and sorting big-first # Splitting video and sorting big-first
timestamps = self.scenedetect(self.args.file_path) timestamps = self.scenedetect(self.args.file_path)
self.split(self.args.file_path, timestamps) self.split(self.args.file_path, timestamps)
files = self.get_video_queue('.temp/split') files = self.get_video_queue(self.temp_dir / 'split')
# Extracting audio # Extracting audio
self.extract_audio(self.args.file_path) self.extract_audio(self.args.file_path)
@ -444,10 +441,10 @@ class Av1an:
for i, _ in enumerate(tqdm(pool.imap_unordered(self.encode, commands), total=len(files), leave=True), 1): for i, _ in enumerate(tqdm(pool.imap_unordered(self.encode, commands), total=len(files), leave=True), 1):
pass pass
self.concatenate_video(self.args.file_path) self.concatenate_video()
# Delete temp folders # Delete temp folders
rmtree(join(self.here, ".temp")) shutil.rmtree(self.temp_dir)
elif self.mode == 1: elif self.mode == 1:
self.image(self.args.file_path) self.image(self.args.file_path)