|
|
@@ -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__':
|
Xintao