480 lines
15 KiB
Python
Executable File
480 lines
15 KiB
Python
Executable File
#!/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", "4500", "-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=35, 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_true", dest="is_10bit", help="Encode using 8-bit profile")
|
|
arger.set_defaults(is_8bit=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()
|