initial commit

This commit is contained in:
2023-02-18 07:39:08 +01:00
commit e63ab7b77e
63 changed files with 7122 additions and 0 deletions
Executable
+479
View File
@@ -0,0 +1,479 @@
#!/usr/bin/env python3.9
import argparse
import dataclasses
import datetime
import json
import logging
import re
import subprocess
import sys
import threading
import time
import typing
from pathlib import Path
import math
import rich.progress
logger = logging.getLogger(__name__)
@dataclasses.dataclass
class FfprobeResult:
duration_sec: float
duration_human: str
codec: str
fps: float
size_bytes: int
width: int
height: int
bitrate: int
container: str
ar: float
sample_ar: float
def ffprobe(video: Path) -> FfprobeResult:
logger.debug("executing ffprobe")
proc = subprocess.run(
args=[
"ffprobe",
"-v",
"error",
"-show_streams",
"-show_format",
"-print_format",
"json",
str(video),
],
text=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
try:
proc.check_returncode()
except subprocess.CalledProcessError:
logger.error("failed to run ffprobe", exc_info=True)
raise
output: dict = json.loads(proc.stdout)
video_stream: dict = [s for s in output.get("streams", []) if s.get("codec_type") == "video"][0]
video_format: dict = output.get("format", {})
width, height = video_stream.get("width"), video_stream.get("height")
codec = video_stream.get("codec_name")
duration = float(video_format.get("duration"))
fps_long = eval(video_stream.get("r_frame_rate", ""))
fps = float(f"{fps_long:.3f}")
size = int(video_format.get("size", 0))
duration_time = datetime.timedelta(seconds=math.ceil(duration))
bitrate = int(video_stream.get("bit_rate", video_format.get("bit_rate", 0)))
try:
sample_w, sample_h = list(map(int, video_stream["sample_aspect_ratio"].split(":")))
sample_ar = (sample_w / sample_h) or 1
except:
sample_ar = 1
logger.debug(f"dimensions={width}x{height} duration={duration_time}, fps={fps}, size={size // 1_048_576}MB")
# total_frames = video_stream.get("nb_frames")
return FfprobeResult(
duration_sec=duration,
duration_human=str(duration_time),
codec=codec,
fps=fps,
size_bytes=size,
width=width,
bitrate=bitrate,
height=height,
container=video.suffix.lstrip(".").lower(),
ar=width / height,
sample_ar=sample_ar,
# total_frames=total_frames,
)
def make_thumbnail_tile(
video: Path,
image_path: Path = None,
columns: int = 3,
interval_seconds: int = 60,
tile_width: int = 540,
skip_first_sec: int = 10,
skip_if_exists: bool = False,
) -> Path:
if not image_path:
image_path = video.parent / f"{video.stem}.thumbnail.jpg"
if image_path.is_file() and skip_if_exists:
logger.debug("thumbnail already exists")
return image_path
info = ffprobe(video)
min_frames = columns * 3
sec_per_frame = min(interval_seconds, math.ceil(info.duration_sec / min_frames))
rows = math.ceil(info.duration_sec // sec_per_frame / columns)
tile = f"{columns}x{rows}"
scaled_width, scaled_height = tile_width, math.ceil(tile_width / info.ar / info.sample_ar)
scale = f"{scaled_width}:{scaled_height}"
proc = subprocess.run(
# fmt: off
args=[
"ffmpeg",
"-v", "error",
"-skip_frame", "nokey",
"-ss", f"{skip_first_sec}",
"-i", str(video.absolute()),
"-vf", f"fps=1/{sec_per_frame},scale={scale},tile={tile}",
"-frames", "1",
"-y", str(image_path.absolute()),
],
# fmt: on
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=2000,
)
proc.check_returncode()
return image_path
@dataclasses.dataclass
class EncodeProgress:
percent: float
fps_avg: int
eta: str
current_size: int
@property
def encoded_mb(self) -> int:
return math.ceil(self.current_size / 1_048_576)
@property
def estimated_mb(self) -> int:
return math.ceil(self.encoded_mb / (self.percent / 100))
@dataclasses.dataclass
class UploadProgress:
uploaded_mb: int
speed: str
percent: float
def noop(*args):
pass
def upload_file(
file_path: Path,
destination: str,
watch: bool = False,
on_progress: typing.Callable[[UploadProgress], None] = noop,
on_resync: typing.Callable[[], None] = noop,
):
size = file_path.stat().st_size
while True:
proc = subprocess.Popen(
["rsync", "--bwlimit", "1100", "-rvPa", str(file_path), destination],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
)
def watch_progress(p: subprocess.Popen):
# 5233864 0% 1.93MB/s 0:21:14
# 12,450,960 30% 2.49MB/s 0:00:11
pattern = re.compile(r"\b(?P<uploaded>[\d,]+)\s+(?P<percent>[\d]+)%\s+(?P<speed>\d+\.\d+[kmg]B/s)\b", re.I)
for line in p.stdout:
if m := pattern.search(line):
parsed = m.groupdict()
progress = UploadProgress(
uploaded_mb=int(parsed["uploaded"].replace(",", "")) // 1_048_576,
speed=parsed["speed"],
percent=float(parsed["percent"]),
)
try:
on_progress(progress)
except:
pass
t = threading.Thread(target=watch_progress, args=(proc,))
t.start()
if proc.wait() != 0:
code = proc.poll()
err = proc.stderr.read()
logger.warning(f"rsync failed with {code=}. error: {err}")
return
if not watch:
return
time.sleep(5)
current_size = file_path.stat().st_size
diff = current_size - size
size = current_size
if not diff:
return
logger.info(f"filesize changed {diff} bytes, re-syncing...")
on_resync()
def encode_video_handbrake(
video_path: Path,
save_path: Path,
quality: int = 30,
from_time: typing.Optional[int] = None,
duration: typing.Optional[int] = None,
on_progress: typing.Callable[[EncodeProgress], None] = noop,
extra_args: typing.List[str] = None,
denoise: bool = False,
is_10bit: bool = True,
):
# fmt: off
args = [
'HandbrakeCLI',
'--format', 'av_mp4',
'--input', str(video_path),
*(['--start-at', f'duration:{from_time}'] if from_time else []),
*(['--stop-at', f'duration:{duration}'] if duration else []),
'--output', str(save_path),
'--optimize',
'--encoder', *['vt_h265_10bit' if is_10bit else 'vt_h265'],
'--quality', str(quality),
'--vfr',
'--aencoder', 'ac3',
'--ab', '160',
'--non-anamorphic',
*(['--hqdn3d', 'light'] if denoise else []),
*(extra_args or []),
# '--json',
'--verbose', '0'
]
# fmt: on
# subprocess.run(args, check=True)
# return
logger.debug("executing handbrake with args: %s", args)
p = subprocess.Popen(args, stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
last_progress: typing.Optional[EncodeProgress] = None
def parse_progress(p: subprocess.Popen):
# Encoding: task 1 of 1, 11.96 % (210.12 fps, avg 206.32 fps, ETA 00h06m14s)
re_progress = re.compile(r"(?P<percent>[\d.]+) % \(.+ avg (?P<fps_avg>[\d.]+) fps, ETA (?P<eta>[^)]+)")
for line in p.stdout:
if m := re_progress.search(line.strip()):
parsed = m.groupdict()
percent = round(float(parsed["percent"]), 1)
fps_avg = math.floor(float(parsed["fps_avg"]))
eta = parsed["eta"]
nonlocal last_progress
last_progress = EncodeProgress(percent, fps_avg, eta, current_size=save_path.stat().st_size)
try:
on_progress(last_progress)
except:
pass
t = threading.Thread(target=parse_progress, args=(p,))
t.start()
try:
if p.wait() != 0:
logger.error(f"handbrake failed: {p.stderr.read()}")
except KeyboardInterrupt:
p.kill()
if last_progress:
logger.info(f"cancelled at {last_progress.percent}%")
logger.info(f"encoded file is {last_progress.encoded_mb}MB, estimated: {last_progress.estimated_mb}MB")
raise
def encode_video_ffmpeg(
video_path: Path,
save_path: Path,
quality: int = 30,
from_time: typing.Optional[int] = None,
duration: typing.Optional[int] = None,
on_progress: typing.Callable[[EncodeProgress], None] = noop,
extra_args: typing.List[str] = None,
is_10bit: bool = True,
**kwargs,
):
# fmt: off
args = [
'ffmpeg',
'-y',
# '-v', 'quiet',
'-progress', 'pipe:1',
'-stats_period', '3',
*(['-ss', str(from_time)] if from_time else []),
*(['-t', str(duration)] if duration else []),
'-i', str(video_path),
'-c:v', 'hevc_videotoolbox',
'-q:v', f'{quality}',
'-profile:v', *['main10' if is_10bit else 'main'],
'-map_metadata', '0',
'-metadata', f'title={video_path.stem}',
*(extra_args or []),
str(save_path),
]
# fmt: on
logger.info(f"calling ffmpeg with {args=}")
proc = subprocess.run(args, check=True)
def parse_args():
arger = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
arger.add_argument("video_path", type=Path, help="path to video file")
arger.add_argument("-e", "--encoder", choices=["handbrake", "ffmpeg"], default="handbrake", help="encoder engine")
arger.add_argument("-q", "--quality", type=float, default=25, help="x265 quality factor")
arger.add_argument("--output-dir", "-o", dest="output_dir", type=Path, help="Dir to save encoded files")
arger.add_argument("--rsync", dest="upload_target", help="rsync encoded file to a host")
arger.add_argument("--validate", action="store_true", default=False, help="perform validations before starting")
arger.add_argument("--denoise", action="store_true", default=False, help="Enable denoise filter (Handbrake only)")
arger.add_argument("--10bit", action="store_true", dest="is_10bit", help="Encode using 10-bit profile")
arger.add_argument("--8bit", action="store_false", dest="is_10bit", help="Encode using 8-bit profile")
arger.set_defaults(is_10bit=True)
def parse_time(val: str) -> int:
parts = val.split(":")
if len(parts) == 1:
return int(val)
elif len(parts) == 2:
return int(parts[0]) * 60 + int(parts[1])
return int(parts[0]) * 60 + int(parts[1]) * 60 + int(parts[2])
arger.add_argument(
"--from",
dest="from_time",
default=0,
type=parse_time,
help="Start encoding from this time. Example 05:00 or 300",
)
arger.add_argument("--duration", type=parse_time, help="Stop encoding at this time. Example 07:00 or 420")
if len(sys.argv[1:]) < 1:
arger.print_help()
exit(1)
args, extra_args = arger.parse_known_args()
if extra_args and extra_args[0] == "--":
extra_args = extra_args[1:]
return args, extra_args
def generate_filename(video_path: Path) -> str:
probe = ffprobe(video_path)
if 1900 <= probe.width <= 2000:
hd = "1080p"
elif 1200 <= probe.width <= 1400:
hd = "720p"
elif 3000 <= probe.width:
hd = "4K"
else:
hd = None
new_stem = re.sub(r"(\[\d+[pk]])", "", video_path.stem)
new_stem = f"{new_stem.strip()} [{hd}, x265]"
return new_stem
def main():
args, extra_args = parse_args()
video_path: Path = args.video_path
video_path = video_path.expanduser().resolve()
if not video_path.is_file():
logger.error("No such file")
exit(1)
if args.validate and len(video_path.name) >= 128:
logger.error("Filename is too long")
exit(1)
save_path = video_path.with_stem(generate_filename(video_path)).with_suffix(".mp4")
if video_path == save_path:
save_path = save_path.with_suffix(".reencoded" + video_path.suffix)
output_dir: Path = args.output_dir
if not output_dir:
output_dir = save_path.parent / "_reenc"
output_dir = output_dir.expanduser().resolve()
output_dir.mkdir(exist_ok=True, parents=True)
save_path = output_dir / save_path.name
bar = rich.progress.Progress(refresh_per_second=2)
task_encode = bar.add_task("encoding", visible=False)
up_task = bar.add_task("uploading", visible=False, start=False)
total_encoded_mb = 0
def sync_later():
logger.info("syncing in 10s")
time.sleep(10)
bar.update(up_task, visible=True)
bar.start_task(up_task)
upload_file(
save_path,
args.upload_target,
watch=True,
on_progress=lambda p: bar.update(
up_task,
description=f"upload: {p.speed}, {p.uploaded_mb:3}MB/{total_encoded_mb:3}MB",
completed=p.uploaded_mb / total_encoded_mb * 100,
),
on_resync=lambda: bar.reset(up_task, description="upload: re-syncing"),
)
t = threading.Thread(target=sync_later)
if args.upload_target:
t.start()
encoder = encode_video_ffmpeg if args.encoder == "ffmpeg" else encode_video_handbrake
with bar:
def _on_encode_progress(p: EncodeProgress):
nonlocal total_encoded_mb
total_encoded_mb = p.encoded_mb
bar.update(
task_encode,
completed=p.percent,
description=f"encode: {p.fps_avg:3}fps, {p.encoded_mb:3}MB/{p.estimated_mb:3}MB",
)
try:
bar.update(task_encode, visible=True)
encoder(
video_path,
quality=args.quality,
save_path=save_path,
extra_args=extra_args,
on_progress=_on_encode_progress,
from_time=args.from_time,
duration=args.duration,
denoise=args.denoise,
is_10bit=args.is_10bit,
)
except KeyboardInterrupt:
logger.info("cancelled")
exit(1)
logger.info(f"finished encoding {video_path}")
if args.upload_target:
t.join()
image_path = make_thumbnail_tile(save_path, image_path=save_path.with_suffix(".jpg"), skip_if_exists=True)
upload_file(image_path, args.upload_target)
if __name__ == "__main__":
logging.basicConfig(level=logging.DEBUG)
main()