304 lines
10 KiB
Python
304 lines
10 KiB
Python
import argparse
|
|
import os
|
|
import queue
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import urllib.request
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
BarColumn,
|
|
MofNCompleteColumn,
|
|
Progress,
|
|
TaskProgressColumn,
|
|
TextColumn,
|
|
TimeElapsedColumn,
|
|
TimeRemainingColumn,
|
|
)
|
|
|
|
console = Console(stderr=True)
|
|
|
|
MODELS = {
|
|
"x4": ("realesrgan-x4plus", 4),
|
|
"x4-anime": ("realesrgan-x4plus-anime", 4),
|
|
"x2": ("realesr-animevideov3-x2", 2),
|
|
"x3": ("realesr-animevideov3-x3", 3),
|
|
"video-x2": ("realesr-animevideov3-x2", 2),
|
|
"video-x3": ("realesr-animevideov3-x3", 3),
|
|
"video-x4": ("realesr-animevideov3-x4", 4),
|
|
}
|
|
|
|
|
|
_RELEASE = "realesrgan-ncnn-vulkan-20220424"
|
|
_DOWNLOAD_URLS = {
|
|
"win32": f"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/{_RELEASE}-windows.zip",
|
|
"linux": f"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/{_RELEASE}-ubuntu.zip",
|
|
"darwin": f"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/{_RELEASE}-macos.zip",
|
|
}
|
|
_EXE_NAME = "realesrgan-ncnn-vulkan.exe" if sys.platform == "win32" else "realesrgan-ncnn-vulkan"
|
|
|
|
|
|
def _download_ncnn(tools: Path):
|
|
platform = sys.platform
|
|
url = _DOWNLOAD_URLS.get(platform)
|
|
if not url:
|
|
console.print(f"[red]no download available for platform: {platform}[/red]")
|
|
sys.exit(1)
|
|
|
|
tools.mkdir(parents=True, exist_ok=True)
|
|
zip_path = tools / f"{_RELEASE}.zip"
|
|
|
|
console.print(f"[dim]downloading {url}[/dim]")
|
|
with Progress(
|
|
TextColumn("{task.description}"),
|
|
BarColumn(),
|
|
TaskProgressColumn(),
|
|
TimeElapsedColumn(),
|
|
TimeRemainingColumn(),
|
|
console=console,
|
|
) as progress:
|
|
task = progress.add_task("Downloading ", total=None)
|
|
|
|
def _reporthook(block: int, block_size: int, total: int):
|
|
if total > 0:
|
|
progress.update(task, total=total, completed=min(block * block_size, total))
|
|
|
|
urllib.request.urlretrieve(url, zip_path, reporthook=_reporthook)
|
|
|
|
console.print("[dim]extracting...[/dim]")
|
|
with zipfile.ZipFile(zip_path, "r") as zf:
|
|
zf.extractall(tools)
|
|
zip_path.unlink()
|
|
|
|
if sys.platform != "win32":
|
|
exe = tools / _RELEASE / _EXE_NAME
|
|
exe.chmod(exe.stat().st_mode | 0o111)
|
|
|
|
|
|
def find_executable() -> Path:
|
|
tools = Path(__file__).parent / ".tools"
|
|
candidates = sorted(tools.glob(f"{_RELEASE}-*"))
|
|
if not candidates:
|
|
_download_ncnn(tools)
|
|
candidates = sorted(tools.glob(f"{_RELEASE}-*"))
|
|
exe = candidates[-1] / _EXE_NAME
|
|
if not exe.exists():
|
|
console.print(f"[red]executable not found: {exe}[/red]")
|
|
sys.exit(1)
|
|
return exe
|
|
|
|
|
|
def _parse_time(t: str) -> float:
|
|
parts = t.split(":")
|
|
return sum(float(p) * 60 ** i for i, p in enumerate(reversed(parts)))
|
|
|
|
|
|
def n_frames_to_seconds(ss: str | None, to: str | None, total: float) -> float:
|
|
start = _parse_time(ss) if ss else 0.0
|
|
end = _parse_time(to) if to else total
|
|
return max(0.0, end - start)
|
|
|
|
|
|
def probe_video(path: str) -> dict:
|
|
result = subprocess.run(
|
|
[
|
|
"ffprobe", "-v", "error", "-select_streams", "v:0",
|
|
"-show_entries", "stream=r_frame_rate,width,height:format=duration",
|
|
"-of", "default=noprint_wrappers=1",
|
|
path,
|
|
],
|
|
capture_output=True, text=True,
|
|
)
|
|
info: dict[str, str] = {}
|
|
for line in result.stdout.splitlines():
|
|
k, _, v = line.partition("=")
|
|
info[k.strip()] = v.strip()
|
|
num, _, den = info.get("r_frame_rate", "30/1").partition("/")
|
|
fps = float(num) / float(den or 1)
|
|
return {
|
|
"fps": fps,
|
|
"duration": float(info.get("duration", 0)),
|
|
"width": int(info.get("width", 0)),
|
|
"height": int(info.get("height", 0)),
|
|
}
|
|
|
|
|
|
def _pipe_reader(stream, q: queue.Queue):
|
|
for line in iter(stream.readline, b""):
|
|
q.put(line.decode(errors="replace").strip())
|
|
q.put(None)
|
|
|
|
|
|
def ffmpeg_with_progress(cmd: list[str], label: str, total: float | None) -> int:
|
|
t = total or 100.0
|
|
with Progress(
|
|
TextColumn("{task.description}"),
|
|
BarColumn(),
|
|
TaskProgressColumn(),
|
|
TimeElapsedColumn(),
|
|
TimeRemainingColumn(),
|
|
console=console,
|
|
) as progress:
|
|
task = progress.add_task(label, total=t)
|
|
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
q: queue.Queue = queue.Queue()
|
|
threading.Thread(target=_pipe_reader, args=(proc.stdout, q), daemon=True).start()
|
|
threading.Thread(target=_pipe_reader, args=(proc.stderr, queue.Queue()), daemon=True).start()
|
|
try:
|
|
while True:
|
|
line = q.get()
|
|
if line is None:
|
|
break
|
|
key, _, value = line.partition("=")
|
|
if key == "out_time" and total:
|
|
try:
|
|
h, m, s = value.split(":")
|
|
elapsed_s = int(h) * 3600 + int(m) * 60 + float(s)
|
|
progress.update(task, completed=min(elapsed_s, t))
|
|
except ValueError:
|
|
pass
|
|
elif key == "progress" and value == "end":
|
|
progress.update(task, completed=t)
|
|
except KeyboardInterrupt:
|
|
proc.terminate()
|
|
proc.wait()
|
|
console.print("\n[yellow]cancelled[/yellow]")
|
|
sys.exit(130)
|
|
proc.wait()
|
|
return proc.returncode
|
|
|
|
|
|
def extract_frames(input_file: str, frames_dir: Path, ss: str | None, to: str | None, duration: float | None) -> int:
|
|
cmd = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y"]
|
|
if ss:
|
|
cmd += ["-ss", ss]
|
|
if to:
|
|
cmd += ["-to", to]
|
|
cmd += ["-i", input_file, str(frames_dir / "%08d.png")]
|
|
cmd += ["-progress", "pipe:1", "-stats_period", "0.5"]
|
|
|
|
rc = ffmpeg_with_progress(cmd, "Extracting ", duration)
|
|
if rc != 0:
|
|
console.print("[red]frame extraction failed[/red]")
|
|
sys.exit(rc)
|
|
|
|
return len(sorted(frames_dir.glob("*.png")))
|
|
|
|
|
|
def upscale_frames(exe: Path, frames_dir: Path, out_dir: Path, model: str, scale: int, gpu_id: int, n_frames: int):
|
|
cmd = [
|
|
str(exe),
|
|
"-i", str(frames_dir),
|
|
"-o", str(out_dir),
|
|
"-n", model,
|
|
"-s", str(scale),
|
|
"-g", str(gpu_id),
|
|
"-f", "png",
|
|
]
|
|
|
|
with Progress(
|
|
TextColumn("{task.description}"),
|
|
BarColumn(),
|
|
MofNCompleteColumn(),
|
|
TimeElapsedColumn(),
|
|
TimeRemainingColumn(),
|
|
console=console,
|
|
) as progress:
|
|
task = progress.add_task("Upscaling ", total=n_frames)
|
|
proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
stop = threading.Event()
|
|
def _watcher():
|
|
while not stop.is_set():
|
|
count = len(list(out_dir.glob("*.png")))
|
|
progress.update(task, completed=count)
|
|
time.sleep(0.5)
|
|
threading.Thread(target=_watcher, daemon=True).start()
|
|
|
|
try:
|
|
proc.wait()
|
|
except KeyboardInterrupt:
|
|
proc.terminate()
|
|
proc.wait()
|
|
stop.set()
|
|
console.print("\n[yellow]cancelled[/yellow]")
|
|
sys.exit(130)
|
|
|
|
stop.set()
|
|
|
|
if proc.returncode != 0:
|
|
console.print("[red]upscaling failed[/red]")
|
|
sys.exit(proc.returncode)
|
|
|
|
|
|
def assemble_video(frames_dir: Path, original: str, output: str, fps: float, clip_duration: float | None, crf: int, ss: str | None, to: str | None):
|
|
audio_input: list[str] = []
|
|
if ss:
|
|
audio_input += ["-ss", ss]
|
|
if to:
|
|
audio_input += ["-to", to]
|
|
audio_input += ["-i", original]
|
|
|
|
cmd = [
|
|
"ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
|
"-framerate", str(fps),
|
|
"-i", str(frames_dir / "%08d.png"),
|
|
*audio_input,
|
|
"-map", "0:v", "-map", "1:a?",
|
|
"-c:v", "libx264", "-crf", str(crf), "-preset", "slow",
|
|
"-c:a", "copy",
|
|
"-progress", "pipe:1", "-stats_period", "0.5",
|
|
output,
|
|
]
|
|
rc = ffmpeg_with_progress(cmd, "Assembling ", clip_duration)
|
|
if rc != 0:
|
|
console.print("[red]assembly failed[/red]")
|
|
sys.exit(rc)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Real-ESRGAN video upscaler (NCNN)")
|
|
parser.add_argument("input", help="Input video file")
|
|
parser.add_argument("output", help="Output video file")
|
|
parser.add_argument("-ss", help="Start time (e.g. 00:01:30 or 90)")
|
|
parser.add_argument("-to", help="End time (e.g. 00:02:00 or 120)")
|
|
parser.add_argument("--model", default="x4", choices=list(MODELS), help="Model (default: x4)")
|
|
parser.add_argument("--gpu", type=int, default=0, help="GPU device ID, -1 for CPU (default: 0)")
|
|
parser.add_argument("--crf", type=int, default=18, help="Output CRF quality (default: 18)")
|
|
|
|
args = parser.parse_args()
|
|
exe = find_executable()
|
|
model_name, scale = MODELS[args.model]
|
|
|
|
info = probe_video(args.input)
|
|
console.print(
|
|
f"[dim]{Path(args.input).name}[/dim] "
|
|
f"[cyan]{info['width']}x{info['height']}[/cyan] "
|
|
f"[cyan]{info['fps']:.2f} fps[/cyan] "
|
|
f"[cyan]{info['duration']:.1f}s[/cyan] "
|
|
f"[dim]→ {model_name} ({scale}x)[/dim]"
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory(prefix="esrgan_in_") as frames_tmp, \
|
|
tempfile.TemporaryDirectory(prefix="esrgan_out_") as out_tmp:
|
|
|
|
frames_dir = Path(frames_tmp)
|
|
out_dir = Path(out_tmp)
|
|
|
|
clip_duration = n_frames_to_seconds(args.ss, args.to, info["duration"])
|
|
n = extract_frames(args.input, frames_dir, args.ss, args.to, clip_duration)
|
|
console.print(f"[dim]{n} frames[/dim]")
|
|
|
|
upscale_frames(exe, frames_dir, out_dir, model_name, scale, args.gpu, n)
|
|
assemble_video(out_dir, args.input, args.output, info["fps"], clip_duration, args.crf, args.ss, args.to)
|
|
|
|
console.print(f"[green]done[/green] {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|