|
import argparse |
|
import cv2 |
|
import glob |
|
import mimetypes |
|
import numpy as np |
|
import os |
|
import shutil |
|
import subprocess |
|
import torch |
|
from basicsr.archs.rrdbnet_arch import RRDBNet |
|
from basicsr.utils.download_util import load_file_from_url |
|
from os import path as osp |
|
from tqdm import tqdm |
|
|
|
from realesrgan import RealESRGANer |
|
from realesrgan.archs.srvgg_arch import SRVGGNetCompact |
|
|
|
try: |
|
import ffmpeg |
|
except ImportError: |
|
import pip |
|
|
|
pip.main(["install", "--user", "ffmpeg-python"]) |
|
import ffmpeg |
|
|
|
|
|
def get_video_meta_info(video_path): |
|
ret = {} |
|
probe = ffmpeg.probe(video_path) |
|
video_streams = [ |
|
stream for stream in probe["streams"] if stream["codec_type"] == "video" |
|
] |
|
has_audio = any(stream["codec_type"] == "audio" for stream in probe["streams"]) |
|
ret["width"] = video_streams[0]["width"] |
|
ret["height"] = video_streams[0]["height"] |
|
ret["fps"] = eval(video_streams[0]["avg_frame_rate"]) |
|
ret["audio"] = ffmpeg.input(video_path).audio if has_audio else None |
|
ret["nb_frames"] = int(video_streams[0]["nb_frames"]) |
|
return ret |
|
|
|
|
|
def get_sub_video(args, num_process, process_idx): |
|
if num_process == 1: |
|
return args.input |
|
meta = get_video_meta_info(args.input) |
|
duration = int(meta["nb_frames"] / meta["fps"]) |
|
part_time = duration // num_process |
|
print(f"duration: {duration}, part_time: {part_time}") |
|
os.makedirs( |
|
osp.join(args.output, f"{args.video_name}_inp_tmp_videos"), exist_ok=True |
|
) |
|
out_path = osp.join( |
|
args.output, f"{args.video_name}_inp_tmp_videos", f"{process_idx:03d}.mp4" |
|
) |
|
cmd = [ |
|
args.ffmpeg_bin, |
|
f"-i {args.input}", |
|
"-ss", |
|
f"{part_time * process_idx}", |
|
f"-to {part_time * (process_idx + 1)}" |
|
if process_idx != num_process - 1 |
|
else "", |
|
"-async 1", |
|
out_path, |
|
"-y", |
|
] |
|
print(" ".join(cmd)) |
|
subprocess.call(" ".join(cmd), shell=True) |
|
return out_path |
|
|
|
|
|
class Reader: |
|
def __init__(self, args, total_workers=1, worker_idx=0): |
|
self.args = args |
|
input_type = mimetypes.guess_type(args.input)[0] |
|
self.input_type = "folder" if input_type is None else input_type |
|
self.paths = [] |
|
self.audio = None |
|
self.input_fps = None |
|
if self.input_type.startswith("video"): |
|
video_path = get_sub_video(args, total_workers, worker_idx) |
|
self.stream_reader = ( |
|
ffmpeg.input(video_path) |
|
.output("pipe:", format="rawvideo", pix_fmt="bgr24", loglevel="error") |
|
.run_async(pipe_stdin=True, pipe_stdout=True, cmd=args.ffmpeg_bin) |
|
) |
|
meta = get_video_meta_info(video_path) |
|
self.width = meta["width"] |
|
self.height = meta["height"] |
|
self.input_fps = meta["fps"] |
|
self.audio = meta["audio"] |
|
self.nb_frames = meta["nb_frames"] |
|
|
|
else: |
|
if self.input_type.startswith("image"): |
|
self.paths = [args.input] |
|
else: |
|
paths = sorted(glob.glob(os.path.join(args.input, "*"))) |
|
tot_frames = len(paths) |
|
num_frame_per_worker = tot_frames // total_workers + ( |
|
1 if tot_frames % total_workers else 0 |
|
) |
|
self.paths = paths[ |
|
num_frame_per_worker |
|
* worker_idx : num_frame_per_worker |
|
* (worker_idx + 1) |
|
] |
|
|
|
self.nb_frames = len(self.paths) |
|
assert self.nb_frames > 0, "empty folder" |
|
from PIL import Image |
|
|
|
tmp_img = Image.open(self.paths[0]) |
|
self.width, self.height = tmp_img.size |
|
self.idx = 0 |
|
|
|
def get_resolution(self): |
|
return self.height, self.width |
|
|
|
def get_fps(self): |
|
if self.args.fps is not None: |
|
return self.args.fps |
|
elif self.input_fps is not None: |
|
return self.input_fps |
|
return 24 |
|
|
|
def get_audio(self): |
|
return self.audio |
|
|
|
def __len__(self): |
|
return self.nb_frames |
|
|
|
def get_frame_from_stream(self): |
|
img_bytes = self.stream_reader.stdout.read( |
|
self.width * self.height * 3 |
|
) |
|
if not img_bytes: |
|
return None |
|
img = np.frombuffer(img_bytes, np.uint8).reshape([self.height, self.width, 3]) |
|
return img |
|
|
|
def get_frame_from_list(self): |
|
if self.idx >= self.nb_frames: |
|
return None |
|
img = cv2.imread(self.paths[self.idx]) |
|
self.idx += 1 |
|
return img |
|
|
|
def get_frame(self): |
|
if self.input_type.startswith("video"): |
|
return self.get_frame_from_stream() |
|
else: |
|
return self.get_frame_from_list() |
|
|
|
def close(self): |
|
if self.input_type.startswith("video"): |
|
self.stream_reader.stdin.close() |
|
self.stream_reader.wait() |
|
|
|
|
|
class Writer: |
|
def __init__(self, args, audio, height, width, video_save_path, fps): |
|
out_width, out_height = int(width * args.outscale), int(height * args.outscale) |
|
if out_height > 2160: |
|
print( |
|
"You are generating video that is larger than 4K, which will be very slow due to IO speed.", |
|
"We highly recommend to decrease the outscale(aka, -s).", |
|
) |
|
|
|
if audio is not None: |
|
self.stream_writer = ( |
|
ffmpeg.input( |
|
"pipe:", |
|
format="rawvideo", |
|
pix_fmt="bgr24", |
|
s=f"{out_width}x{out_height}", |
|
framerate=fps, |
|
) |
|
.output( |
|
audio, |
|
video_save_path, |
|
pix_fmt="yuv420p", |
|
vcodec="libx264", |
|
loglevel="error", |
|
acodec="copy", |
|
) |
|
.overwrite_output() |
|
.run_async(pipe_stdin=True, pipe_stdout=True, cmd=args.ffmpeg_bin) |
|
) |
|
else: |
|
self.stream_writer = ( |
|
ffmpeg.input( |
|
"pipe:", |
|
format="rawvideo", |
|
pix_fmt="bgr24", |
|
s=f"{out_width}x{out_height}", |
|
framerate=fps, |
|
) |
|
.output( |
|
video_save_path, |
|
pix_fmt="yuv420p", |
|
vcodec="libx264", |
|
loglevel="error", |
|
) |
|
.overwrite_output() |
|
.run_async(pipe_stdin=True, pipe_stdout=True, cmd=args.ffmpeg_bin) |
|
) |
|
|
|
def write_frame(self, frame): |
|
frame = frame.astype(np.uint8).tobytes() |
|
self.stream_writer.stdin.write(frame) |
|
|
|
def close(self): |
|
self.stream_writer.stdin.close() |
|
self.stream_writer.wait() |
|
|
|
|
|
def inference_video(args, video_save_path, device=None, total_workers=1, worker_idx=0): |
|
|
|
args.model_name = args.model_name.split(".pth")[0] |
|
if args.model_name == "RealESRGAN_x4plus": |
|
model = RRDBNet( |
|
num_in_ch=3, |
|
num_out_ch=3, |
|
num_feat=64, |
|
num_block=23, |
|
num_grow_ch=32, |
|
scale=4, |
|
) |
|
netscale = 4 |
|
file_url = [ |
|
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth" |
|
] |
|
elif args.model_name == "RealESRNet_x4plus": |
|
model = RRDBNet( |
|
num_in_ch=3, |
|
num_out_ch=3, |
|
num_feat=64, |
|
num_block=23, |
|
num_grow_ch=32, |
|
scale=4, |
|
) |
|
netscale = 4 |
|
file_url = [ |
|
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth" |
|
] |
|
elif ( |
|
args.model_name == "RealESRGAN_x4plus_anime_6B" |
|
): |
|
model = RRDBNet( |
|
num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4 |
|
) |
|
netscale = 4 |
|
file_url = [ |
|
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth" |
|
] |
|
elif args.model_name == "RealESRGAN_x2plus": |
|
model = RRDBNet( |
|
num_in_ch=3, |
|
num_out_ch=3, |
|
num_feat=64, |
|
num_block=23, |
|
num_grow_ch=32, |
|
scale=2, |
|
) |
|
netscale = 2 |
|
file_url = [ |
|
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth" |
|
] |
|
elif args.model_name == "realesr-animevideov3": |
|
model = SRVGGNetCompact( |
|
num_in_ch=3, |
|
num_out_ch=3, |
|
num_feat=64, |
|
num_conv=16, |
|
upscale=4, |
|
act_type="prelu", |
|
) |
|
netscale = 4 |
|
file_url = [ |
|
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth" |
|
] |
|
elif args.model_name == "realesr-general-x4v3": |
|
model = SRVGGNetCompact( |
|
num_in_ch=3, |
|
num_out_ch=3, |
|
num_feat=64, |
|
num_conv=32, |
|
upscale=4, |
|
act_type="prelu", |
|
) |
|
netscale = 4 |
|
file_url = [ |
|
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth", |
|
"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", |
|
] |
|
|
|
|
|
model_path = os.path.join("weights", args.model_name + ".pth") |
|
if not os.path.isfile(model_path): |
|
ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) |
|
for url in file_url: |
|
|
|
model_path = load_file_from_url( |
|
url=url, |
|
model_dir=os.path.join(ROOT_DIR, "weights"), |
|
progress=True, |
|
file_name=None, |
|
) |
|
|
|
|
|
dni_weight = None |
|
if args.model_name == "realesr-general-x4v3" and args.denoise_strength != 1: |
|
wdn_model_path = model_path.replace( |
|
"realesr-general-x4v3", "realesr-general-wdn-x4v3" |
|
) |
|
model_path = [model_path, wdn_model_path] |
|
dni_weight = [args.denoise_strength, 1 - args.denoise_strength] |
|
|
|
|
|
upsampler = RealESRGANer( |
|
scale=netscale, |
|
model_path=model_path, |
|
dni_weight=dni_weight, |
|
model=model, |
|
tile=args.tile, |
|
tile_pad=args.tile_pad, |
|
pre_pad=args.pre_pad, |
|
half=not args.fp32, |
|
device=device, |
|
) |
|
|
|
if "anime" in args.model_name and args.face_enhance: |
|
print( |
|
"face_enhance is not supported in anime models, we turned this option off for you. " |
|
"if you insist on turning it on, please manually comment the relevant lines of code." |
|
) |
|
args.face_enhance = False |
|
|
|
if args.face_enhance: |
|
from gfpgan import GFPGANer |
|
|
|
face_enhancer = GFPGANer( |
|
model_path="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth", |
|
upscale=args.outscale, |
|
arch="clean", |
|
channel_multiplier=2, |
|
bg_upsampler=upsampler, |
|
) |
|
else: |
|
face_enhancer = None |
|
|
|
reader = Reader(args, total_workers, worker_idx) |
|
audio = reader.get_audio() |
|
height, width = reader.get_resolution() |
|
fps = reader.get_fps() |
|
writer = Writer(args, audio, height, width, video_save_path, fps) |
|
|
|
pbar = tqdm(total=len(reader), unit="frame", desc="inference") |
|
while True: |
|
img = reader.get_frame() |
|
if img is None: |
|
break |
|
|
|
try: |
|
if args.face_enhance: |
|
_, _, output = face_enhancer.enhance( |
|
img, has_aligned=False, only_center_face=False, paste_back=True |
|
) |
|
else: |
|
output, _ = upsampler.enhance(img, outscale=args.outscale) |
|
except RuntimeError as error: |
|
print("Error", error) |
|
print( |
|
"If you encounter CUDA out of memory, try to set --tile with a smaller number." |
|
) |
|
else: |
|
writer.write_frame(output) |
|
|
|
torch.cuda.synchronize(device) |
|
pbar.update(1) |
|
|
|
reader.close() |
|
writer.close() |
|
|
|
|
|
def run(args): |
|
args.video_name = osp.splitext(os.path.basename(args.input))[0] |
|
video_save_path = osp.join(args.output, f"{args.video_name}_{args.suffix}.mp4") |
|
|
|
if args.extract_frame_first: |
|
tmp_frames_folder = osp.join(args.output, f"{args.video_name}_inp_tmp_frames") |
|
os.makedirs(tmp_frames_folder, exist_ok=True) |
|
os.system( |
|
f"ffmpeg -i {args.input} -qscale:v 1 -qmin 1 -qmax 1 -vsync 0 {tmp_frames_folder}/frame%08d.png" |
|
) |
|
args.input = tmp_frames_folder |
|
|
|
num_gpus = torch.cuda.device_count() |
|
num_process = num_gpus * args.num_process_per_gpu |
|
if num_process == 1: |
|
inference_video(args, video_save_path) |
|
return |
|
|
|
ctx = torch.multiprocessing.get_context("spawn") |
|
pool = ctx.Pool(num_process) |
|
os.makedirs( |
|
osp.join(args.output, f"{args.video_name}_out_tmp_videos"), exist_ok=True |
|
) |
|
pbar = tqdm(total=num_process, unit="sub_video", desc="inference") |
|
for i in range(num_process): |
|
sub_video_save_path = osp.join( |
|
args.output, f"{args.video_name}_out_tmp_videos", f"{i:03d}.mp4" |
|
) |
|
pool.apply_async( |
|
inference_video, |
|
args=( |
|
args, |
|
sub_video_save_path, |
|
torch.device(i % num_gpus), |
|
num_process, |
|
i, |
|
), |
|
callback=lambda arg: pbar.update(1), |
|
) |
|
pool.close() |
|
pool.join() |
|
|
|
|
|
|
|
with open(f"{args.output}/{args.video_name}_vidlist.txt", "w") as f: |
|
for i in range(num_process): |
|
f.write(f"file '{args.video_name}_out_tmp_videos/{i:03d}.mp4'\n") |
|
|
|
cmd = [ |
|
args.ffmpeg_bin, |
|
"-f", |
|
"concat", |
|
"-safe", |
|
"0", |
|
"-i", |
|
f"{args.output}/{args.video_name}_vidlist.txt", |
|
"-c", |
|
"copy", |
|
f"{video_save_path}", |
|
] |
|
print(" ".join(cmd)) |
|
subprocess.call(cmd) |
|
shutil.rmtree(osp.join(args.output, f"{args.video_name}_out_tmp_videos")) |
|
if osp.exists(osp.join(args.output, f"{args.video_name}_inp_tmp_videos")): |
|
shutil.rmtree(osp.join(args.output, f"{args.video_name}_inp_tmp_videos")) |
|
os.remove(f"{args.output}/{args.video_name}_vidlist.txt") |
|
|
|
|
|
def main(): |
|
"""Inference demo for Real-ESRGAN. |
|
It mainly for restoring anime videos. |
|
|
|
""" |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument( |
|
"-i", "--input", type=str, default="inputs", help="Input video, image or folder" |
|
) |
|
parser.add_argument( |
|
"-n", |
|
"--model_name", |
|
type=str, |
|
default="realesr-animevideov3", |
|
help=( |
|
"Model names: realesr-animevideov3 | RealESRGAN_x4plus_anime_6B | RealESRGAN_x4plus | RealESRNet_x4plus |" |
|
" RealESRGAN_x2plus | realesr-general-x4v3" |
|
"Default:realesr-animevideov3" |
|
), |
|
) |
|
parser.add_argument( |
|
"-o", "--output", type=str, default="results", help="Output folder" |
|
) |
|
parser.add_argument( |
|
"-dn", |
|
"--denoise_strength", |
|
type=float, |
|
default=0.5, |
|
help=( |
|
"Denoise strength. 0 for weak denoise (keep noise), 1 for strong denoise ability. " |
|
"Only used for the realesr-general-x4v3 model" |
|
), |
|
) |
|
parser.add_argument( |
|
"-s", |
|
"--outscale", |
|
type=float, |
|
default=4, |
|
help="The final upsampling scale of the image", |
|
) |
|
parser.add_argument( |
|
"--suffix", type=str, default="out", help="Suffix of the restored video" |
|
) |
|
parser.add_argument( |
|
"-t", |
|
"--tile", |
|
type=int, |
|
default=0, |
|
help="Tile size, 0 for no tile during testing", |
|
) |
|
parser.add_argument("--tile_pad", type=int, default=10, help="Tile padding") |
|
parser.add_argument( |
|
"--pre_pad", type=int, default=0, help="Pre padding size at each border" |
|
) |
|
parser.add_argument( |
|
"--face_enhance", action="store_true", help="Use GFPGAN to enhance face" |
|
) |
|
parser.add_argument( |
|
"--fp32", |
|
action="store_true", |
|
help="Use fp32 precision during inference. Default: fp16 (half precision).", |
|
) |
|
parser.add_argument( |
|
"--fps", type=float, default=None, help="FPS of the output video" |
|
) |
|
parser.add_argument( |
|
"--ffmpeg_bin", type=str, default="ffmpeg", help="The path to ffmpeg" |
|
) |
|
parser.add_argument("--extract_frame_first", action="store_true") |
|
parser.add_argument("--num_process_per_gpu", type=int, default=1) |
|
|
|
parser.add_argument( |
|
"--alpha_upsampler", |
|
type=str, |
|
default="realesrgan", |
|
help="The upsampler for the alpha channels. Options: realesrgan | bicubic", |
|
) |
|
parser.add_argument( |
|
"--ext", |
|
type=str, |
|
default="auto", |
|
help="Image extension. Options: auto | jpg | png, auto means using the same extension as inputs", |
|
) |
|
args = parser.parse_args() |
|
|
|
args.input = args.input.rstrip("/").rstrip("\\") |
|
os.makedirs(args.output, exist_ok=True) |
|
|
|
if mimetypes.guess_type(args.input)[0] is not None and mimetypes.guess_type( |
|
args.input |
|
)[0].startswith("video"): |
|
is_video = True |
|
else: |
|
is_video = False |
|
|
|
if is_video and args.input.endswith(".flv"): |
|
mp4_path = args.input.replace(".flv", ".mp4") |
|
os.system(f"ffmpeg -i {args.input} -codec copy {mp4_path}") |
|
args.input = mp4_path |
|
|
|
if args.extract_frame_first and not is_video: |
|
args.extract_frame_first = False |
|
|
|
run(args) |
|
|
|
if args.extract_frame_first: |
|
tmp_frames_folder = osp.join(args.output, f"{args.video_name}_inp_tmp_frames") |
|
shutil.rmtree(tmp_frames_folder) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
