OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2024-12-29 16:09:10 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #103 from sat3ll/wip_1

Browse Source 
use os.path.join
This commit is contained in:
K4YT3X 2019-06-15 14:26:28 -04:00 committed by GitHub
commit 76144b589a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 41 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
bin/ffmpeg.py
View File

@ -12,6 +12,7 @@ operations.
from avalon_framework import Avalon from avalon_framework import Avalon
import json import json
import subprocess import subprocess
import os
class Ffmpeg: class Ffmpeg:
@ -26,13 +27,9 @@ class Ffmpeg:
self.ffmpeg_settings = ffmpeg_settings self.ffmpeg_settings = ffmpeg_settings
self.ffmpeg_path = self.ffmpeg_settings['ffmpeg_path'] self.ffmpeg_path = self.ffmpeg_settings['ffmpeg_path']
# add a forward slash to directory if not present
# otherwise there will be a format error
if self.ffmpeg_path[-1] != '/' and self.ffmpeg_path[-1] != '\\':
self.ffmpeg_path = f'{self.ffmpeg_path}\\'
self.ffmpeg_binary = f'{self.ffmpeg_path}ffmpeg.exe' self.ffmpeg_binary = os.path.join(self.ffmpeg_path, 'ffmpeg.exe')
self.ffmpeg_probe_binary = f'{self.ffmpeg_path}ffprobe.exe' self.ffmpeg_probe_binary = os.path.join(self.ffmpeg_path, 'ffprobe.exe')
self.image_format = image_format self.image_format = image_format
def get_video_info(self, input_video): def get_video_info(self, input_video):
@ -88,7 +85,7 @@ class Ffmpeg:
execute.extend(self._read_configuration(phase='video_to_frames', section='output_options')) execute.extend(self._read_configuration(phase='video_to_frames', section='output_options'))
execute.extend([ execute.extend([
f'{extracted_frames}\\extracted_%0d.{self.image_format}' os.path.join(extracted_frames, f'extracted_%0d.{self.image_format}')
]) ])
execute.extend(self._read_configuration(phase='video_to_frames')) execute.extend(self._read_configuration(phase='video_to_frames'))
@ -120,7 +117,7 @@ class Ffmpeg:
# append input frames path into command # append input frames path into command
execute.extend([ execute.extend([
'-i', '-i',
f'{upscaled_frames}\\extracted_%d.{self.image_format}' os.path.join(upscaled_frames, f'extracted_%d.{self.image_format}')
]) ])
# read FFmpeg output options # read FFmpeg output options
@ -131,7 +128,7 @@ class Ffmpeg:
# specify output file location # specify output file location
execute.extend([ execute.extend([
f'{upscaled_frames}\\no_audio.mp4' os.path.join(upscaled_frames, 'no_audio.mp4')
]) ])
self._execute(execute) self._execute(execute)
@ -147,7 +144,7 @@ class Ffmpeg:
execute = [ execute = [
self.ffmpeg_binary, self.ffmpeg_binary,
'-i', '-i',
f'{upscaled_frames}\\no_audio.mp4', os.path.join(upscaled_frames, 'no_audio.mp4'),
'-i', '-i',
input_video input_video
] ]

8
bin/upscaler.py
View File

@ -53,7 +53,7 @@ class Upscaler:
self.scale_ratio = None self.scale_ratio = None
self.model_dir = None self.model_dir = None
self.threads = 5 self.threads = 5
self.video2x_cache_directory = f'{tempfile.gettempdir()}\\video2x' self.video2x_cache_directory = os.path.join(tempfile.gettempdir(), 'video2x')
self.image_format = 'png' self.image_format = 'png'
self.preserve_frames = False self.preserve_frames = False
@ -144,7 +144,7 @@ class Upscaler:
# if this thread is not empty, then an exception has occured # if this thread is not empty, then an exception has occured
self.upscaler_exceptions = [] self.upscaler_exceptions = []
# initialize waifu2x driver # initialize waifu2x driver
if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter': if self.waifu2x_driver != 'waifu2x_caffe' and self.waifu2x_driver != 'waifu2x_converter':
raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}') raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
@ -159,7 +159,7 @@ class Upscaler:
w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.image_format, self.upscaler_exceptions) w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.image_format, self.upscaler_exceptions)
for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]: for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]:
renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image) renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image)
shutil.move(f'{self.upscaled_frames}\\{image}', f'{self.upscaled_frames}\\{renamed}') shutil.move(os.path.join(self.upscaled_frames, image), os.path.join(self.upscaled_frames, renamed))
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
progress_bar.join() progress_bar.join()
@ -182,7 +182,7 @@ class Upscaler:
thread_pool = [] thread_pool = []
thread_directories = [] thread_directories = []
for thread_id in range(self.threads): for thread_id in range(self.threads):
thread_directory = f'{self.extracted_frames}\\{str(thread_id)}' thread_directory = os.path.join(self.extracted_frames, str(thread_id))
thread_directories.append(thread_directory) thread_directories.append(thread_directory)
# delete old directories and create new directories # delete old directories and create new directories

20
bin/video2x.py
View File

@ -80,7 +80,7 @@ def process_arguments():
upscaler_options.add_argument('-d', '--driver', help='waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter']) upscaler_options.add_argument('-d', '--driver', help='waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter'])
upscaler_options.add_argument('-y', '--model_dir', help='directory containing model JSON files', action='store') upscaler_options.add_argument('-y', '--model_dir', help='directory containing model JSON files', action='store')
upscaler_options.add_argument('-t', '--threads', help='number of threads to use for upscaling', action='store', type=int, default=1) upscaler_options.add_argument('-t', '--threads', help='number of threads to use for upscaling', action='store', type=int, default=1)
upscaler_options.add_argument('-c', '--config', help='video2x config file location', action='store', default=f'{os.path.dirname(os.path.abspath(sys.argv[0]))}\\video2x.json') upscaler_options.add_argument('-c', '--config', help='video2x config file location', action='store', default=os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), 'video2x.json'))
upscaler_options.add_argument('-b', '--batch', help='enable batch mode (select all default values to questions)', action='store_true') upscaler_options.add_argument('-b', '--batch', help='enable batch mode (select all default values to questions)', action='store_true')
# scaling options # scaling options
@ -123,7 +123,8 @@ def check_memory():
# check if Nvidia-smi is available # check if Nvidia-smi is available
# GPUtil requires nvidia-smi.exe to interact with GPU # GPUtil requires nvidia-smi.exe to interact with GPU
if args.method == 'gpu' or args.method == 'cudnn': if args.method == 'gpu' or args.method == 'cudnn':
if not os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe'): if not (shutil.which('nvidia-smi') or
os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe')):
# Nvidia System Management Interface not available # Nvidia System Management Interface not available
Avalon.warning('Nvidia-smi not available, skipping available memory check') Avalon.warning('Nvidia-smi not available, skipping available memory check')
Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using') Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using')
@ -191,20 +192,23 @@ def absolutify_paths(config):
# check waifu2x-caffe path # check waifu2x-caffe path
if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE): if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE):
config['waifu2x_caffe']['waifu2x_caffe_path'] = f'{current_directory}\\{config["waifu2x_caffe"]["waifu2x_caffe_path"]}' config['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(current_directory,
config['waifu2x_caffe']['waifu2x_caffe_path'])
# check waifu2x-converter-cpp path # check waifu2x-converter-cpp path
if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE): if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE):
config['waifu2x_converter']['waifu2x_converter_path'] = f'{current_directory}\\{config["waifu2x_converter"]["waifu2x_converter_path"]}' config['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(current_directory,
config['waifu2x_converter']['waifu2x_converter_path'])
# check ffmpeg path # check ffmpeg path
if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE): if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE):
config['ffmpeg']['ffmpeg_path'] = f'{current_directory}\\{config["ffmpeg"]["ffmpeg_path"]}' config['ffmpeg']['ffmpeg_path'] = os.path.join(current_directory, config['ffmpeg']['ffmpeg_path'])
# check video2x cache path # check video2x cache path
if config['video2x']['video2x_cache_directory']: if config['video2x']['video2x_cache_directory']:
if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE): if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE):
config['video2x']['video2x_cache_directory'] = f'{current_directory}\\{config["video2x"]["video2x_cache_directory"]}' config['video2x']['video2x_cache_directory'] = os.path.join(current_directory,
config['video2x']['video2x_cache_directory'])
return config return config
@ -279,7 +283,7 @@ preserve_frames = config['video2x']['preserve_frames']
# create temp directories if they don't exist # create temp directories if they don't exist
if not video2x_cache_directory: if not video2x_cache_directory:
video2x_cache_directory = f'{tempfile.gettempdir()}\\video2x' video2x_cache_directory = os.path.join(tempfile.gettempdir(), 'video2x')
if video2x_cache_directory and not os.path.isdir(video2x_cache_directory): if video2x_cache_directory and not os.path.isdir(video2x_cache_directory):
if not os.path.isfile(video2x_cache_directory) and not os.path.islink(video2x_cache_directory): if not os.path.isfile(video2x_cache_directory) and not os.path.islink(video2x_cache_directory):
@ -341,7 +345,7 @@ try:
# upscale videos in a directory # upscale videos in a directory
Avalon.info(f'Upscaling videos in directory: {args.input}') Avalon.info(f'Upscaling videos in directory: {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 = f'{args.output}\\{input_video}' output_video = os.path.join(args.output, input_video)
upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings) upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings)
# set optional options # set optional options

23
bin/video2x_setup.py
View File

@ -110,7 +110,7 @@ class Video2xSetup:
self.trash.append(ffmpeg_zip) self.trash.append(ffmpeg_zip)
with zipfile.ZipFile(ffmpeg_zip) as zipf: with zipfile.ZipFile(ffmpeg_zip) as zipf:
zipf.extractall(f'{os.getenv("localappdata")}\\video2x') zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x'))
def _install_waifu2x_caffe(self): def _install_waifu2x_caffe(self):
""" Install waifu2x_caffe """ Install waifu2x_caffe
@ -127,7 +127,7 @@ class Video2xSetup:
self.trash.append(waifu2x_caffe_zip) self.trash.append(waifu2x_caffe_zip)
with zipfile.ZipFile(waifu2x_caffe_zip) as zipf: with zipfile.ZipFile(waifu2x_caffe_zip) as zipf:
zipf.extractall(f'{os.getenv("localappdata")}\\video2x') zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x'))
def _install_waifu2x_converter_cpp(self): def _install_waifu2x_converter_cpp(self):
""" Install waifu2x_caffe """ Install waifu2x_caffe
@ -145,7 +145,7 @@ class Video2xSetup:
self.trash.append(waifu2x_converter_cpp_zip) self.trash.append(waifu2x_converter_cpp_zip)
with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf: with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf:
zipf.extractall(f'{os.getenv("localappdata")}\\video2x\\waifu2x-converter-cpp') zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-converter-cpp'))
def _generate_config(self): def _generate_config(self):
""" Generate video2x config """ Generate video2x config
@ -159,14 +159,19 @@ class Video2xSetup:
# configure only the specified drivers # configure only the specified drivers
if self.driver == 'all': if self.driver == 'all':
template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = f'{local_app_data}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe' template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x',
template_dict['waifu2x_converter']['waifu2x_converter_path'] = f'{local_app_data}\\video2x\\waifu2x-converter-cpp' 'waifu2x-caffe', 'waifu2x-caffe-cui.exe')
template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x',
'waifu2x-converter-cpp')
elif self.driver == 'waifu2x_caffe': elif self.driver == 'waifu2x_caffe':
template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = f'{local_app_data}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe' template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x',
'waifu2x-caffe', 'waifu2x-caffe-cui.exe')
elif self.driver == 'waifu2x_converter': elif self.driver == 'waifu2x_converter':
template_dict['waifu2x_converter']['waifu2x_converter_path'] = f'{local_app_data}\\video2x\\waifu2x-converter-cpp' template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x',
'waifu2x-converter-cpp')
template_dict['ffmpeg']['ffmpeg_path'] = f'{local_app_data}\\video2x\\ffmpeg-latest-win64-static\\bin' template_dict['ffmpeg']['ffmpeg_path'] = os.path.join(local_app_data, 'video2x',
'ffmpeg-latest-win64-static', 'bin')
template_dict['video2x']['video2x_cache_directory'] = None template_dict['video2x']['video2x_cache_directory'] = None
template_dict['video2x']['preserve_frames'] = False template_dict['video2x']['preserve_frames'] = False
@ -182,7 +187,7 @@ def download(url, save_path, chunk_size=4096):
from tqdm import tqdm from tqdm import tqdm
import requests import requests
output_file = f'{save_path}\\{url.split("/")[-1]}' output_file = os.path.join(save_path, url.split("/")[-1])
print(f'Downloading: {url}') print(f'Downloading: {url}')
print(f'Chunk size: {chunk_size}') print(f'Chunk size: {chunk_size}')
print(f'Saving to: {output_file}') print(f'Saving to: {output_file}')

6
bin/waifu2x_converter.py
View File

@ -12,6 +12,7 @@ for waifu2x-converter-cpp.
from avalon_framework import Avalon from avalon_framework import Avalon
import subprocess import subprocess
import threading import threading
import os
class Waifu2xConverter: class Waifu2xConverter:
@ -59,7 +60,8 @@ class Waifu2xConverter:
# models_rgb must be specified manually for waifu2x-converter-cpp # models_rgb must be specified manually for waifu2x-converter-cpp
# if it's not specified in the arguments, create automatically # if it's not specified in the arguments, create automatically
if self.waifu2x_settings['model-dir'] is None: if self.waifu2x_settings['model-dir'] is None:
self.waifu2x_settings['model-dir'] = f'{self.waifu2x_settings["waifu2x_converter_path"]}\\models_rgb' self.waifu2x_settings['model-dir'] = os.path.join(self.waifu2x_settings['waifu2x_converter_path'],
'models_rgb')
# print thread start message # print thread start message
self.print_lock.acquire() self.print_lock.acquire()
@ -75,7 +77,7 @@ class Waifu2xConverter:
# the key doesn't need to be passed in this case # the key doesn't need to be passed in this case
if key == 'waifu2x_converter_path': if key == 'waifu2x_converter_path':
execute.append(f'{str(value)}\\waifu2x-converter-cpp.exe') execute.append(os.path.join(str(value), 'waifu2x-converter-cpp.exe'))
# null or None means that leave this option out (keep default) # null or None means that leave this option out (keep default)
elif value is None or value is False: elif value is None or value is False: