2.5.0 added progress bar, changed default cache folder location

This commit is contained in:
k4yt3x 2019-03-04 19:35:50 -05:00
parent 1f8dc0dd47
commit d5de0740b4

View File

@ -13,7 +13,7 @@ __ __ _ _ ___ __ __
Name: Video2X Controller Name: Video2X Controller
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: February 26, 2019 Last Modified: March 4, 2019
Licensed under the GNU General Public License Version 3 (GNU GPL v3), Licensed under the GNU General Public License Version 3 (GNU GPL v3),
available at: https://www.gnu.org/licenses/gpl-3.0.txt available at: https://www.gnu.org/licenses/gpl-3.0.txt
@ -28,15 +28,17 @@ details or quality, keeping lines smooth and edges sharp.
from avalon_framework import Avalon from avalon_framework import Avalon
from upscaler import Upscaler from upscaler import Upscaler
from upscaler import MODELS_AVAILABLE from upscaler import MODELS_AVAILABLE
import argparse
import GPUtil import GPUtil
import argparse
import json import json
import os import os
import psutil import psutil
import shutil
import tempfile
import time import time
import traceback import traceback
VERSION = '2.4.3' VERSION = '2.5.0'
# Each thread might take up to 2.5 GB during initialization. # Each thread might take up to 2.5 GB during initialization.
# (system memory, not to be confused with GPU memory) # (system memory, not to be confused with GPU memory)
@ -181,24 +183,25 @@ waifu2x_path = config['waifu2x_path']
ffmpeg_path = config['ffmpeg_path'] ffmpeg_path = config['ffmpeg_path']
ffmpeg_arguments = config['ffmpeg_arguments'] ffmpeg_arguments = config['ffmpeg_arguments']
ffmpeg_hwaccel = config['ffmpeg_hwaccel'] ffmpeg_hwaccel = config['ffmpeg_hwaccel']
extracted_frames = config['extracted_frames'] video2x_cache_folder = config['video2x_cache_folder']
upscaled_frames = config['upscaled_frames']
preserve_frames = config['preserve_frames'] preserve_frames = config['preserve_frames']
# Create temp directories if they don't exist # Create temp directories if they don't exist
for directory in [extracted_frames, upscaled_frames]: if not video2x_cache_folder:
if directory and not os.path.isdir(directory): video2x_cache_folder = '{}\\video2x'.format(tempfile.gettempdir())
if not os.path.isfile(directory) and not os.path.islink(directory):
Avalon.warning('Specified temporary folder/directory {} does not exist'.format(directory)) if video2x_cache_folder and not os.path.isdir(video2x_cache_folder):
if Avalon.ask('Create folder/directory?'): if not os.path.isfile(video2x_cache_folder) and not os.path.islink(video2x_cache_folder):
if os.mkdir(directory) == 0: Avalon.warning('Specified cache folder/directory {} does not exist'.format(video2x_cache_folder))
Avalon.info('{} created'.format(directory)) if Avalon.ask('Create folder/directory?', True):
if os.mkdir(video2x_cache_folder) is None:
Avalon.info('{} created'.format(video2x_cache_folder))
else: else:
Avalon.error('Unable to create {}'.format(directory)) Avalon.error('Unable to create {}'.format(video2x_cache_folder))
Avalon.error('Aborting...') Avalon.error('Aborting...')
exit(1) exit(1)
else: else:
Avalon.error('Specified temporary folder/directory is a file/link') Avalon.error('Specified cache folder/directory is a file/link')
Avalon.error('Unable to continue, exiting...') Avalon.error('Unable to continue, exiting...')
exit(1) exit(1)
@ -211,20 +214,29 @@ try:
if os.path.isfile(args.input): if os.path.isfile(args.input):
""" Upscale single video file """ """ Upscale single video file """
Avalon.info('Upscaling single video file: {}'.format(args.input)) Avalon.info('Upscaling single video file: {}'.format(args.input))
upscaler = Upscaler(input_video=args.input, output_video=args.output, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, waifu2x_driver=args.driver, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, ratio=args.ratio, model_type=args.model_type, threads=args.threads, extracted_frames=extracted_frames, upscaled_frames=upscaled_frames) upscaler = Upscaler(input_video=args.input, output_video=args.output, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, waifu2x_driver=args.driver, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, ratio=args.ratio, model_type=args.model_type, threads=args.threads, video2x_cache_folder=video2x_cache_folder)
upscaler.run() upscaler.run()
upscaler.cleanup()
elif os.path.isdir(args.input): elif os.path.isdir(args.input):
""" Upscale videos in a folder/directory """ """ Upscale videos in a folder/directory """
Avalon.info('Upscaling videos in folder/directory: {}'.format(args.input)) Avalon.info('Upscaling videos in folder/directory: {}'.format(args.input))
for input_video in [f for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f))]: for input_video in [f for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f))]:
output_video = '{}\\{}'.format(args.output, input_video) output_video = '{}\\{}'.format(args.output, input_video)
upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, waifu2x_driver=args.driver, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, ratio=args.ratio, model_type=args.model_type, threads=args.threads, extracted_frames=extracted_frames, upscaled_frames=upscaled_frames) upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, waifu2x_driver=args.driver, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, ratio=args.ratio, model_type=args.model_type, threads=args.threads, video2x_cache_folder=video2x_cache_folder)
upscaler.run() upscaler.run()
upscaler.cleanup()
else: else:
Avalon.error('Input path is neither a file nor a folder/directory') Avalon.error('Input path is neither a file nor a folder/directory')
raise FileNotFoundError('{} is neither file nor folder/directory'.format(args.input)) raise FileNotFoundError('{} is neither file nor folder/directory'.format(args.input))
Avalon.info('Program completed, taking {} seconds'.format(round((time.time() - begin_time), 5))) Avalon.info('Program completed, taking {} seconds'.format(round((time.time() - begin_time), 5)))
except Exception: except Exception:
Avalon.error('An exception occurred') Avalon.error('An exception has occurred')
traceback.print_exc() traceback.print_exc()
finally:
# Remove Video2X Cache folder
try:
if not preserve_frames:
shutil.rmtree(video2x_cache_folder)
except FileNotFoundError:
pass