lucasrufkahr/Real-ESRGAN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

support realesr-animevideov3

Browse Source
This commit is contained in:
Xintao
2022-04-24 17:22:43 +08:00
parent 9ff1944d06
commit 98add035f2
2 changed files with 46 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
inference_realesrgan_video.py
View File

@@ -24,11 +24,10 @@ def main():
'-n',
'--model_name',
type=str,
default='RealESRGAN_x4plus',
help=('Model names: RealESRGAN_x4plus | RealESRNet_x4plus | RealESRGAN_x4plus_anime_6B | RealESRGAN_x2plus | '
'RealESRGANv2-anime-xsx2 | RealESRGANv2-animevideo-xsx2-nousm | RealESRGANv2-animevideo-xsx2 | '
'RealESRGANv2-anime-xsx4 | RealESRGANv2-animevideo-xsx4-nousm | RealESRGANv2-animevideo-xsx4. '
'Default: RealESRGAN_x4plus'))
default='realesr-animevideov3',
help=('Model names: realesr-animevideov3 | RealESRGAN_x4plus_anime_6B | RealESRGAN_x4plus | RealESRNet_x4plus |'
' RealESRGAN_x2plus | '
'Default:realesr-animevideov3'))
parser.add_argument('-o', '--output', type=str, default='results', help='Output folder')
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')
@@ -36,9 +35,8 @@ def main():
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('--half', action='store_true', help='Use half precision during inference')
parser.add_argument('-v', '--video', action='store_true', help='Output a video using ffmpeg')
parser.add_argument('-a', '--audio', action='store_true', help='Keep audio')
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('--consumer', type=int, default=4, help='Number of IO consumers')
@@ -55,7 +53,7 @@ def main():
args = parser.parse_args()
# ---------------------- determine models according to model names ---------------------- #
args.model_name = args.model_name.split('.')[0]
args.model_name = args.model_name.split('.pth')[0]
if args.model_name in ['RealESRGAN_x4plus', 'RealESRNet_x4plus']: # x4 RRDBNet model
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
netscale = 4
@@ -65,14 +63,7 @@ def main():
elif args.model_name in ['RealESRGAN_x2plus']: # x2 RRDBNet model
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
netscale = 2
elif args.model_name in [
'RealESRGANv2-anime-xsx2', 'RealESRGANv2-animevideo-xsx2-nousm', 'RealESRGANv2-animevideo-xsx2'
]: # x2 VGG-style model (XS size)
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
netscale = 2
elif args.model_name in [
'RealESRGANv2-anime-xsx4', 'RealESRGANv2-animevideo-xsx4-nousm', 'RealESRGANv2-animevideo-xsx4'
]: # x4 VGG-style model (XS size)
elif args.model_name in ['realesr-animevideov3']: # x4 VGG-style model (XS size)
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
netscale = 4
@@ -91,7 +82,7 @@ def main():
tile=args.tile,
tile_pad=args.tile_pad,
pre_pad=args.pre_pad,
half=args.half)
half=not args.fp32)
if args.face_enhance: # Use GFPGAN for face enhancement
from gfpgan import GFPGANer
@@ -106,6 +97,7 @@ def main():
save_frame_folder = os.path.join(args.output, 'frames_tmpout')
os.makedirs(save_frame_folder, exist_ok=True)
# input can be a video file / a folder of frames / an image
if mimetypes.guess_type(args.input)[0].startswith('video'): # is a video file
video_name = os.path.splitext(os.path.basename(args.input))[0]
frame_folder = os.path.join('tmp_frames', video_name)
@@ -114,13 +106,13 @@ def main():
os.system(f'ffmpeg -i {args.input} -qscale:v 1 -qmin 1 -qmax 1 -vsync 0 {frame_folder}/frame%08d.png')
# get image path list
paths = sorted(glob.glob(os.path.join(frame_folder, '*')))
if args.video:
if args.fps is None:
# get input video fps
import ffmpeg
probe = ffmpeg.probe(args.input)
video_streams = [stream for stream in probe['streams'] if stream['codec_type'] == 'video']
args.fps = eval(video_streams[0]['avg_frame_rate'])
# get input video fps
if args.fps is None:
import ffmpeg
probe = ffmpeg.probe(args.input)
video_streams = [stream for stream in probe['streams'] if stream['codec_type'] == 'video']
args.fps = eval(video_streams[0]['avg_frame_rate'])
elif mimetypes.guess_type(args.input)[0].startswith('image'): # is an image file
paths = [args.input]
video_name = 'video'
@@ -180,20 +172,13 @@ def main():
pbar.close()
# merge frames to video
if args.video:
video_save_path = os.path.join(args.output, f'{video_name}_{args.suffix}.mp4')
if args.audio:
os.system(
f'ffmpeg -r {args.fps} -i {save_frame_folder}/frame%08d_out.{extension} -i {args.input}'
f' -map 0:v:0 -map 1:a:0 -c:a copy -c:v libx264 -r {args.fps} -pix_fmt yuv420p {video_save_path}')
else:
os.system(f'ffmpeg -r {args.fps} -i {save_frame_folder}/frame%08d_out.{extension} '
f'-c:v libx264 -r {args.fps} -pix_fmt yuv420p {video_save_path}')
# delete tmp file
shutil.rmtree(save_frame_folder)
if os.path.isdir(frame_folder):
shutil.rmtree(frame_folder)
video_save_path = os.path.join(args.output, f'{video_name}_{args.suffix}.mp4')
os.system(f'ffmpeg -r {args.fps} -i {save_frame_folder}/frame%08d_out.{extension} -i {args.input}'
f' -map 0:v:0 -map 1:a:0 -c:a copy -c:v libx264 -r {args.fps} -pix_fmt yuv420p {video_save_path}')
# delete tmp file
shutil.rmtree(save_frame_folder)
if os.path.isdir(frame_folder):
shutil.rmtree(frame_folder)
if __name__ == '__main__':