while True:
fileName = Avalon.gets('请输入要保存的文件名，必须以.md为后缀:')
if fileName.split('.')[-1] != 'md':
Avalon.warning('文件名错误！')
else:
try:
with open(fileName, 'w+', encoding='utf-8') as f:
if not f.writable:
raise Exception
except:
Avalon.warning('文件错误!')
else:
break
Avalon.debug_info('程序已经启动...正在获取帖子信息')
post = api(postID=postID, seeLZ=onlySeeLZ, debug=GENERAL_DEBUG_MODE)
postInfo = post.getInfo()
Avalon.time_info('开始任务:"%s"(作者:%s)' % (postInfo['Title'], postInfo['Author']),
highlight=True)
for i in range(1, postInfo['TotalPage'] + 1):
Avalon.time_info('开始第%d页,共%d页' % (i, postInfo['TotalPage']))
try:
pageHTMLContent = post.getContent(i)
pageMarkdownContent = post.contentToMarkdown(pageHTMLContent,
useImageBed=USE_IMAGE_BED)
except KeyboardInterrupt:
Avalon.critical('用户强制退出')
quit(1)
else:
post.saveToFile(fileName, pageMarkdownContent)
def upscale(self, input_directory, output_directory, scale_ratio, upscaler_exceptions):
"""This is the core function for WAIFU2X class
Arguments:
input_directory {string} -- source directory path
output_directory {string} -- output directory path
ratio {int} -- output video ratio
"""
try:
# overwrite config file settings
self.waifu2x_settings['input'] = input_directory
self.waifu2x_settings['output'] = output_directory
# print thread start message
self.print_lock.acquire()
Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} started')
self.print_lock.release()
# waifu2x_ncnn_vulkan does not have long-opts, we'll have a dictionary that maps "our" config long-opt
# names to their short opts
waifu2x_ncnn_vulkan_opt_flag = {
'input': '-i',
'output': '-o',
'noise-level': '-n',
'scale-ratio': '-s',
'tile-size': '-t',
'model-path': '-m',
'gpu': '-g',
'load-proc-save_threads': '-j',
'verbose': '-v'
}
execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']]
for key in self.waifu2x_settings.keys():
value = self.waifu2x_settings[key]
if key == 'waifu2x_ncnn_vulkan_path':
continue
elif key == 'input':
execute.append(waifu2x_ncnn_vulkan_opt_flag[key])
execute.append(input_directory)
elif key == 'output':
execute.append(waifu2x_ncnn_vulkan_opt_flag[key])
execute.append(output_directory)
elif key == 'scale-ratio':
execute.append(waifu2x_ncnn_vulkan_opt_flag[key])
# waifu2x_ncnn_vulkan does not accept an arbitrary scale ratio, max is 2
if scale_ratio == 1:
execute.append('1')
else:
execute.append('2')
# allow upper if cases to take precedence
elif value is None or value is False:
continue
else:
execute.append(waifu2x_ncnn_vulkan_opt_flag[key])
execute.append(str(value))
Avalon.debug_info(f'Executing: {execute}')
subprocess.run(execute, check=True, stderr=subprocess.DEVNULL)
# print thread exiting message
self.print_lock.acquire()
Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} exiting')
self.print_lock.release()
return 0
except Exception as e:
upscaler_exceptions.append(e)
def run(self):
"""Main controller for Video2X
This function controls the flow of video conversion
and handles all necessary functions.
"""
# external stop signal when called in a thread
self.running = True
# define process pool to contain processes
self.process_pool = []
# load driver modules
DriverWrapperMain = getattr(
importlib.import_module(f"wrappers.{self.driver}"), "WrapperMain")
self.driver_object = DriverWrapperMain(self.driver_settings)
# load options from upscaler class into driver settings
self.driver_object.load_configurations(self)
# initialize FFmpeg object
self.ffmpeg_object = Ffmpeg(
self.ffmpeg_settings,
extracted_frame_format=self.extracted_frame_format)
# define processing queue
self.processing_queue = queue.Queue()
Avalon.info(_("Loading files into processing queue"))
Avalon.debug_info(_("Input path(s): {}").format(self.input))
# make output directory if the input is a list or a directory
if isinstance(self.input, list) or self.input.is_dir():
self.output.mkdir(parents=True, exist_ok=True)
input_files = []
# if input is single directory
# put it in a list for compability with the following code
if not isinstance(self.input, list):
input_paths = [self.input]
else:
input_paths = self.input
# flatten directories into file paths
for input_path in input_paths:
# if the input path is a single file
# add the file's path object to input_files
if input_path.is_file():
input_files.append(input_path)
# if the input path is a directory
# add all files under the directory into the input_files (non-recursive)
elif input_path.is_dir():
input_files.extend(
[f for f in input_path.iterdir() if f.is_file()])
output_paths = []
for input_path in input_files:
# get file type
# try python-magic if it's available
try:
input_file_mime_type = magic.from_file(str(
input_path.absolute()),
mime=True)
input_file_type = input_file_mime_type.split("/")[0]
input_file_subtype = input_file_mime_type.split("/")[1]
except Exception:
input_file_mime_type = input_file_type = input_file_subtype = ""
# if python-magic doesn't determine the file to be an image/video file
# fall back to mimetypes to guess the file type based on the extension
if input_file_type not in ["image", "video"]:
# in case python-magic fails to detect file type
# try guessing file mime type with mimetypes
input_file_mime_type = mimetypes.guess_type(input_path.name)[0]
input_file_type = input_file_mime_type.split("/")[0]
input_file_subtype = input_file_mime_type.split("/")[1]
Avalon.debug_info(
_("File MIME type: {}").format(input_file_mime_type))
# set default output file suffixes
# if image type is GIF, default output suffix is also .gif
if input_file_mime_type == "image/gif":
output_path = self.output / self.output_file_name_format_string.format(
original_file_name=input_path.stem, extension=".gif")
elif input_file_type == "image":
output_path = self.output / self.output_file_name_format_string.format(
original_file_name=input_path.stem,
extension=self.image_output_extension,
)
elif input_file_type == "video":
output_path = self.output / self.output_file_name_format_string.format(
original_file_name=input_path.stem,
extension=self.video_output_extension,
)
# if file is none of: image, image/gif, video
# skip to the next task
else:
Avalon.error(
_("File {} ({}) neither an image nor a video").format(
input_path, input_file_mime_type))
Avalon.warning(_("Skipping this file"))
continue
# if there is only one input file
# do not modify output file suffix
if isinstance(self.input, pathlib.Path) and self.input.is_file():
output_path = self.output
output_path_id = 0
while str(output_path) in output_paths:
output_path = output_path.parent / pathlib.Path(
f"{output_path.stem}_{output_path_id}{output_path.suffix}")
output_path_id += 1
# record output path
output_paths.append(str(output_path))
# push file information into processing queue
self.processing_queue.put((
input_path.absolute(),
output_path.absolute(),
input_file_mime_type,
input_file_type,
input_file_subtype,
))
# check argument sanity before running
self._check_arguments()
# record file count for external calls
self.total_files = self.processing_queue.qsize()
Avalon.info(_("Loaded files into processing queue"))
# print all files in queue for debugging
for job in self.processing_queue.queue:
Avalon.debug_info(_("Input file: {}").format(job[0].absolute()))
try:
while not self.processing_queue.empty():
# get new job from queue
(
self.current_input_file,
output_path,
input_file_mime_type,
input_file_type,
input_file_subtype,
) = self.processing_queue.get()
# get current job starting time for GUI calculations
self.current_processing_starting_time = time.time()
# get video information JSON using FFprobe
Avalon.info(_("Reading file information"))
file_info = self.ffmpeg_object.probe_file_info(
self.current_input_file)
# create temporary directories for storing frames
self.create_temp_directories()
# start handling input
# if input file is a static image
if input_file_type == "image" and input_file_subtype != "gif":
Avalon.info(_("Starting upscaling image"))
# copy original file into the pre-processing directory
shutil.copy(
self.current_input_file,
self.extracted_frames / self.current_input_file.name,
)
width = int(file_info["streams"][0]["width"])
height = int(file_info["streams"][0]["height"])
framerate = self.total_frames = 1
# elif input_file_mime_type == 'image/gif' or input_file_type == 'video':
else:
Avalon.info(_("Starting upscaling video/GIF"))
# find index of video stream
video_stream_index = None
for stream in file_info["streams"]:
if stream["codec_type"] == "video":
video_stream_index = stream["index"]
break
# exit if no video stream found
if video_stream_index is None:
Avalon.error(_("Aborting: No video stream found"))
raise StreamNotFoundError("no video stream found")
# get average frame rate of video stream
framerate = float(
Fraction(file_info["streams"][video_stream_index]
["r_frame_rate"]))
width = int(
file_info["streams"][video_stream_index]["width"])
height = int(
file_info["streams"][video_stream_index]["height"])
# get total number of frames
Avalon.info(
_("Getting total number of frames in the file"))
# if container stores total number of frames in nb_frames, fetch it directly
if "nb_frames" in file_info["streams"][video_stream_index]:
self.total_frames = int(
file_info["streams"][video_stream_index]
["nb_frames"])
# otherwise call FFprobe to count the total number of frames
else:
self.total_frames = self.ffmpeg_object.get_number_of_frames(
self.current_input_file, video_stream_index)
# calculate scale width/height/ratio and scaling jobs if required
Avalon.info(_("Calculating scaling parameters"))
# create a local copy of the global output settings
output_scale = self.scale_ratio
output_width = self.scale_width
output_height = self.scale_height
# calculate output width and height if scale ratio is specified
if output_scale is not None:
output_width = int(
math.ceil(width * output_scale / 2.0) * 2)
output_height = int(
math.ceil(height * output_scale / 2.0) * 2)
else:
# scale keeping aspect ratio is only one of width/height is given
if output_width == 0 or output_width is None:
output_width = output_height / height * width
elif output_height == 0 or output_height is None:
output_height = output_width / width * height
output_width = int(math.ceil(output_width / 2.0) * 2)
output_height = int(math.ceil(output_height / 2.0) * 2)
# calculate required minimum scale ratio
output_scale = max(output_width / width,
output_height / height)
# if driver is one of the drivers that doesn't support arbitrary scaling ratio
# TODO: more documentations on this block
if self.driver in DRIVER_FIXED_SCALING_RATIOS:
# select the optimal driver scaling ratio to use
supported_scaling_ratios = sorted(
DRIVER_FIXED_SCALING_RATIOS[self.driver])
remaining_scaling_ratio = math.ceil(output_scale)
self.scaling_jobs = []
# if the scaling ratio is 1.0
# apply the smallest scaling ratio available
if remaining_scaling_ratio == 1:
self.scaling_jobs.append(supported_scaling_ratios[0])
else:
while remaining_scaling_ratio > 1:
for ratio in supported_scaling_ratios:
if ratio >= remaining_scaling_ratio:
self.scaling_jobs.append(ratio)
remaining_scaling_ratio /= ratio
break
else:
found = False
for i in supported_scaling_ratios:
for j in supported_scaling_ratios:
if i * j >= remaining_scaling_ratio:
self.scaling_jobs.extend([i, j])
remaining_scaling_ratio /= i * j
found = True
break
if found is True:
break
if found is False:
self.scaling_jobs.append(
supported_scaling_ratios[-1])
remaining_scaling_ratio /= supported_scaling_ratios[
-1]
else:
self.scaling_jobs = [output_scale]
# print file information
Avalon.debug_info(_("Framerate: {}").format(framerate))
Avalon.debug_info(_("Width: {}").format(width))
Avalon.debug_info(_("Height: {}").format(height))
Avalon.debug_info(
_("Total number of frames: {}").format(self.total_frames))
Avalon.debug_info(_("Output width: {}").format(output_width))
Avalon.debug_info(_("Output height: {}").format(output_height))
Avalon.debug_info(
_("Required scale ratio: {}").format(output_scale))
Avalon.debug_info(
_("Upscaling jobs queue: {}").format(self.scaling_jobs))
# extract frames from video
if input_file_mime_type == "image/gif" or input_file_type == "video":
self.process_pool.append(
(self.ffmpeg_object.extract_frames(
self.current_input_file, self.extracted_frames)))
self._wait()
# if driver is waifu2x-caffe
# pass pixel format output depth information
if self.driver == "waifu2x_caffe":
# get a dict of all pixel formats and corresponding bit depth
pixel_formats = self.ffmpeg_object.get_pixel_formats()
# try getting pixel format's corresponding bti depth
try:
self.driver_settings["output_depth"] = pixel_formats[
self.ffmpeg_object.pixel_format]
except KeyError:
Avalon.error(
_("Unsupported pixel format: {}").format(
self.ffmpeg_object.pixel_format))
raise UnsupportedPixelError(
f"unsupported pixel format {self.ffmpeg_object.pixel_format}"
)
# upscale images one by one using waifu2x
Avalon.info(_("Starting to upscale extracted frames"))
upscale_begin_time = time.time()
self.current_pass = 1
if self.driver == "waifu2x_caffe":
self.driver_object.set_scale_resolution(
output_width, output_height)
else:
self.driver_object.set_scale_ratio(self.scaling_jobs[0])
self._upscale_frames(self.extracted_frames,
self.upscaled_frames)
for job in self.scaling_jobs[1:]:
self.current_pass += 1
self.driver_object.set_scale_ratio(job)
shutil.rmtree(self.extracted_frames)
shutil.move(self.upscaled_frames, self.extracted_frames)
self.upscaled_frames.mkdir(parents=True, exist_ok=True)
self._upscale_frames(self.extracted_frames,
self.upscaled_frames)
Avalon.info(_("Upscaling completed"))
Avalon.info(
_("Average processing speed: {} seconds per frame").format(
self.total_frames /
(time.time() - upscale_begin_time)))
# downscale frames with Lanczos
Avalon.info(_("Lanczos downscaling frames"))
shutil.rmtree(self.extracted_frames)
shutil.move(self.upscaled_frames, self.extracted_frames)
self.upscaled_frames.mkdir(parents=True, exist_ok=True)
for image in tqdm(
[
i for i in self.extracted_frames.iterdir()
if i.is_file()
and i.name.endswith(self.extracted_frame_format)
],
ascii=True,
desc=_("Downscaling"),
):
image_object = Image.open(image)
# if the image dimensions are not equal to the output size
# resize the image using Lanczos
if (image_object.width, image_object.height) != (
output_width,
output_height,
):
image_object.resize(
(output_width, output_height), Image.LANCZOS).save(
self.upscaled_frames / image.name)
image_object.close()
# if the image's dimensions are already equal to the output size
# move image to the finished directory
else:
image_object.close()
shutil.move(image, self.upscaled_frames / image.name)
# start handling output
# output can be either GIF or video
if input_file_type == "image" and input_file_subtype != "gif":
Avalon.info(_("Exporting image"))
# there should be only one image in the directory
shutil.move(
[
f for f in self.upscaled_frames.iterdir()
if f.is_file()
][0],
output_path,
)
# elif input_file_mime_type == 'image/gif' or input_file_type == 'video':
else:
# if the desired output is gif file
if output_path.suffix.lower() == ".gif":
Avalon.info(
_("Converting extracted frames into GIF image"))
gifski_object = Gifski(self.gifski_settings)
self.process_pool.append(
gifski_object.make_gif(
self.upscaled_frames,
output_path,
framerate,
self.extracted_frame_format,
output_width,
output_height,
))
self._wait()
Avalon.info(_("Conversion completed"))
# if the desired output is video
else:
# frames to video
Avalon.info(
_("Converting extracted frames into video"))
self.process_pool.append(
self.ffmpeg_object.assemble_video(
framerate, self.upscaled_frames))
# f'{scale_width}x{scale_height}'
self._wait()
Avalon.info(_("Conversion completed"))
try:
# migrate audio tracks and subtitles
Avalon.info(
_("Migrating audio, subtitles and other streams to upscaled video"
))
self.process_pool.append(
self.ffmpeg_object.migrate_streams(
self.current_input_file,
output_path,
self.upscaled_frames,
))
self._wait()
# if failed to copy streams
# use file with only video stream
except subprocess.CalledProcessError:
traceback.print_exc()
Avalon.error(_("Failed to migrate streams"))
Avalon.warning(
_("Trying to output video without additional streams"
))
if input_file_mime_type == "image/gif":
# copy will overwrite destination content if exists
shutil.copy(
self.upscaled_frames /
self.ffmpeg_object.intermediate_file_name,
output_path,
)
else:
# construct output file path
output_file_name = f"{output_path.stem}{self.ffmpeg_object.intermediate_file_name.suffix}"
output_video_path = (output_path.parent /
output_file_name)
# if output file already exists
# create temporary directory in output folder
# temporary directories generated by tempfile are guaranteed to be unique
# and won't conflict with other files
if output_video_path.exists():
Avalon.error(_("Output video file exists"))
temporary_directory = pathlib.Path(
tempfile.mkdtemp(
dir=output_path.parent))
output_video_path = (temporary_directory /
output_file_name)
Avalon.info(
_("Created temporary directory to contain file"
))
# move file to new destination
Avalon.info(
_("Writing intermediate file to: {}").
format(output_video_path.absolute()))
shutil.move(
self.upscaled_frames /
self.ffmpeg_object.intermediate_file_name,
output_video_path,
)
# increment total number of files processed
self.cleanup_temp_directories()
self.processing_queue.task_done()
self.total_processed += 1
except (Exception, KeyboardInterrupt, SystemExit) as e:
with contextlib.suppress(ValueError, AttributeError):
self.cleanup_temp_directories()
self.running = False
raise e
# signal upscaling completion
self.running = False
# resolve NTP server hostname into IP address
# this makes sure that the hostname/IP is valid
try:
ntp_server_ip = str(netaddr.IPAddress(socket.gethostbyname(ntp_server)))
except (socket.gaierror, netaddr.core.AddrFormatError) as error:
Avalon.error(f"Unresolvable or invalid NTP hostname: {ntp_server}")
raise error
thread = threading.Thread(target=attack, args=(ntp_server_ip, args.target, args.count,))
thread.daemon = True
thread.name = str(ntp_id)
thread_pool.append(thread)
Avalon.info('Launching attack')
Avalon.info('Press Ctrl+C to stop the script')
# start all threads after they have been added
for thread in thread_pool:
Avalon.debug_info(f"Starting thread {thread.name}")
thread.start()
# sleep and wait for SystemExit and KeyboardInterrupt
try:
for thread in thread_pool:
thread.join()
except (SystemExit, KeyboardInterrupt):
Avalon.warning("Exiting signal received, script exiting")
finally:
Avalon.info("Script finished")
def generate_configs(output_path):
""" Generate configuration file for every peer
This function reads the PEERS list, generates a
configuration file for every peer, and export into
the CONFIG_OUTPUT directory.
"""
if len(pm.peers) == 0:
Avalon.warning('No peers configured, exiting')
exit(0)
if len(pm.peers) == 1:
Avalon.warning('Only one peer configured')
Avalon.info('Generating configuration files')
# Abort is destination is a file / link
if os.path.isfile(output_path) or os.path.islink(output_path):
Avalon.warning('Destination path is a file / link')
Avalon.warning('Aborting configuration generation')
return 1
# Ask if user wants to create the output directory if it doesn't exist
if not os.path.isdir(output_path):
if Avalon.ask(
'Output directory doesn\'t exist. Create output directory?',
True):
os.mkdir(output_path)
else:
Avalon.warning('Aborting configuration generation')
return 1
# Iterate through all peers and generate configuration for each peer
for peer in pm.peers:
Avalon.debug_info('Generating configuration file for {}'.format(
peer.address))
with open('{}/{}.conf'.format(output_path,
peer.address.split('/')[0]),
'w') as config:
# Write Interface configuration
config.write('[Interface]\n')
config.write('PrivateKey = {}\n'.format(peer.private_key))
if peer.address != '':
config.write('Address = {}\n'.format(peer.address))
if peer.listen_port != '':
config.write('ListenPort = {}\n'.format(peer.listen_port))
# Write peers' information
for p in pm.peers:
if p.address == peer.address:
# Skip if peer is self
continue
config.write('\n[Peer]\n')
print(p.private_key)
config.write('PublicKey = {}\n'.format(wg.pubkey(
p.private_key)))
config.write('AllowedIPs = {}\n'.format(p.address))
if p.public_address != '':
config.write('Endpoint = {}:{}\n'.format(
p.public_address, p.listen_port))
if peer.keep_alive:
config.write('PersistentKeepalive = 25\n')
if p.preshared_key:
config.write('PresharedKey = {}\n'.format(p.preshared_key))
def upscale(self, input_directory, output_directory, scale_ratio,
scale_width, scale_height, image_format, upscaler_exceptions):
"""This is the core function for WAIFU2X class
Arguments:
input_directory {string} -- source directory path
output_directory {string} -- output directory path
width {int} -- output video width
height {int} -- output video height
"""
try:
# overwrite config file settings
self.waifu2x_settings['input_path'] = input_directory
self.waifu2x_settings['output_path'] = output_directory
if scale_ratio:
self.waifu2x_settings['scale_ratio'] = scale_ratio
elif scale_width and scale_height:
self.waifu2x_settings['scale_width'] = scale_width
self.waifu2x_settings['scale_height'] = scale_height
self.waifu2x_settings['output_extention'] = image_format
# print thread start message
self.print_lock.acquire()
Avalon.debug_info(
f'[upscaler] Thread {threading.current_thread().name} started')
self.print_lock.release()
# list to be executed
execute = []
execute.append(self.waifu2x_settings['waifu2x_caffe_path'])
for key in self.waifu2x_settings.keys():
value = self.waifu2x_settings[key]
# is executable key or null or None means that leave this option out (keep default)
if key == 'waifu2x_caffe_path' or value is None or value is False:
continue
else:
if len(key) == 1:
execute.append(f'-{key}')
else:
execute.append(f'--{key}')
execute.append(str(value))
Avalon.debug_info(f'Executing: {execute}')
completed_command = subprocess.run(execute, check=True)
# print thread exiting message
self.print_lock.acquire()
Avalon.debug_info(
f'[upscaler] Thread {threading.current_thread().name} exiting')
self.print_lock.release()
# return command execution return code
return completed_command.returncode
except Exception as e:
upscaler_exceptions.append(e)
def _upscale_frames(self):
""" Upscale video frames with waifu2x-caffe
This function upscales all the frames extracted
by ffmpeg using the waifu2x-caffe binary.
Arguments:
w2 {Waifu2x Object} -- initialized waifu2x object
"""
# progress bar process exit signal
self.progress_bar_exit_signal = False
# initialize waifu2x driver
if self.waifu2x_driver not in AVAILABLE_DRIVERS:
raise UnrecognizedDriverError(
f'Unrecognized driver: {self.waifu2x_driver}')
# create a container for all upscaler processes
upscaler_processes = []
# list all images in the extracted frames
frames = [(self.extracted_frames / f)
for f in self.extracted_frames.iterdir() if f.is_file]
# if we have less images than processes,
# create only the processes necessary
if len(frames) < self.processes:
self.processes = len(frames)
# create a directory for each process and append directory
# name into a list
process_directories = []
for process_id in range(self.processes):
process_directory = self.extracted_frames / str(process_id)
process_directories.append(process_directory)
# delete old directories and create new directories
if process_directory.is_dir():
shutil.rmtree(process_directory)
process_directory.mkdir(parents=True, exist_ok=True)
# waifu2x-converter-cpp will perform multi-threading within its own process
if self.waifu2x_driver in ['waifu2x_converter', 'anime4k']:
process_directories = [self.extracted_frames]
else:
# evenly distribute images into each directory
# until there is none left in the directory
for image in frames:
# move image
image.rename(process_directories[0] / image.name)
# rotate list
process_directories = process_directories[
-1:] + process_directories[:-1]
# create threads and start them
for process_directory in process_directories:
# if the driver being used is waifu2x-caffe
if self.waifu2x_driver == 'waifu2x_caffe':
driver = Waifu2xCaffe(copy.deepcopy(self.driver_settings),
self.method, self.model_dir,
self.bit_depth)
if self.scale_ratio:
upscaler_processes.append(
driver.upscale(process_directory, self.upscaled_frames,
self.scale_ratio, False, False,
self.image_format))
else:
upscaler_processes.append(
driver.upscale(process_directory, self.upscaled_frames,
False, self.scale_width,
self.scale_height, self.image_format))
# if the driver being used is waifu2x-converter-cpp
elif self.waifu2x_driver == 'waifu2x_converter':
driver = Waifu2xConverter(self.driver_settings, self.model_dir)
upscaler_processes.append(
driver.upscale(process_directory, self.upscaled_frames,
self.scale_ratio, self.processes,
self.image_format))
# if the driver being used is waifu2x-ncnn-vulkan
elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan':
driver = Waifu2xNcnnVulkan(copy.deepcopy(self.driver_settings))
upscaler_processes.append(
driver.upscale(process_directory, self.upscaled_frames,
self.scale_ratio))
# if the driver being used is anime4k
elif self.waifu2x_driver == 'anime4k':
driver = Anime4k(copy.deepcopy(self.driver_settings))
upscaler_processes += driver.upscale(process_directory,
self.upscaled_frames,
self.scale_ratio,
self.processes)
# if the driver being used is srmd_ncnn_vulkan
elif self.waifu2x_driver == 'srmd_ncnn_vulkan':
driver = SrmdNcnnVulkan(copy.deepcopy(self.driver_settings))
upscaler_processes.append(
driver.upscale(process_directory, self.upscaled_frames,
self.scale_ratio))
# start progress bar in a different thread
progress_bar = threading.Thread(target=self._progress_bar,
args=(process_directories, ))
progress_bar.start()
# create the clearer and start it
Avalon.debug_info('Starting upscaled image cleaner')
image_cleaner = ImageCleaner(self.extracted_frames,
self.upscaled_frames,
len(upscaler_processes))
image_cleaner.start()
# wait for all process to exit
try:
Avalon.debug_info('Main process waiting for subprocesses to exit')
for process in upscaler_processes:
Avalon.debug_info(
f'Subprocess {process.pid} exited with code {process.wait()}'
)
except (KeyboardInterrupt, SystemExit):
Avalon.warning('Exit signal received')
Avalon.warning('Killing processes')
for process in upscaler_processes:
process.terminate()
# cleanup and exit with exit code 1
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
sys.exit(1)
# if the driver is waifu2x-converter-cpp
# images need to be renamed to be recognizable for FFmpeg
if self.waifu2x_driver == 'waifu2x_converter':
for image in [
f for f in self.upscaled_frames.iterdir() if f.is_file()
]:
renamed = re.sub(
f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.image_format}',
f'.{self.image_format}', str(image.name))
(self.upscaled_frames / image).rename(self.upscaled_frames /
renamed)
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
# pass exit signal to progress bar thread
self.progress_bar_exit_signal = True
def _upscale_frames(self):
""" Upscale video frames with waifu2x-caffe
This function upscales all the frames extracted
by ffmpeg using the waifu2x-caffe binary.
Arguments:
w2 {Waifu2x Object} -- initialized waifu2x object
"""
# initialize waifu2x driver
if self.driver not in AVAILABLE_DRIVERS:
raise UnrecognizedDriverError(
_('Unrecognized driver: {}').format(self.driver))
# list all images in the extracted frames
frames = [(self.extracted_frames / f)
for f in self.extracted_frames.iterdir() if f.is_file]
# if we have less images than processes,
# create only the processes necessary
if len(frames) < self.processes:
self.processes = len(frames)
# create a directory for each process and append directory
# name into a list
process_directories = []
for process_id in range(self.processes):
process_directory = self.extracted_frames / str(process_id)
process_directories.append(process_directory)
# delete old directories and create new directories
if process_directory.is_dir():
shutil.rmtree(process_directory)
process_directory.mkdir(parents=True, exist_ok=True)
# waifu2x-converter-cpp will perform multi-threading within its own process
if self.driver in [
'waifu2x_converter_cpp', 'waifu2x_ncnn_vulkan',
'srmd_ncnn_vulkan', 'realsr_ncnn_vulkan', 'anime4kcpp'
]:
process_directories = [self.extracted_frames]
else:
# evenly distribute images into each directory
# until there is none left in the directory
for image in frames:
# move image
image.rename(process_directories[0] / image.name)
# rotate list
process_directories = process_directories[
-1:] + process_directories[:-1]
# create driver processes and start them
for process_directory in process_directories:
self.process_pool.append(
self.driver_object.upscale(process_directory,
self.upscaled_frames))
# start progress bar in a different thread
Avalon.debug_info(_('Starting progress monitor'))
self.progress_monitor = ProgressMonitor(self, process_directories)
self.progress_monitor.start()
# create the clearer and start it
Avalon.debug_info(_('Starting upscaled image cleaner'))
self.image_cleaner = ImageCleaner(self.extracted_frames,
self.upscaled_frames,
len(self.process_pool))
self.image_cleaner.start()
# wait for all process to exit
try:
self._wait()
except (Exception, KeyboardInterrupt, SystemExit) as e:
# cleanup
Avalon.debug_info(_('Killing progress monitor'))
self.progress_monitor.stop()
Avalon.debug_info(_('Killing upscaled image cleaner'))
self.image_cleaner.stop()
raise e
# if the driver is waifu2x-converter-cpp
# images need to be renamed to be recognizable for FFmpeg
if self.driver == 'waifu2x_converter_cpp':
for image in [
f for f in self.upscaled_frames.iterdir() if f.is_file()
]:
renamed = re.sub(
f'_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.extracted_frame_format}',
f'.{self.extracted_frame_format}', str(image.name))
(self.upscaled_frames / image).rename(self.upscaled_frames /
renamed)
# upscaling done, kill helper threads
Avalon.debug_info(_('Killing progress monitor'))
self.progress_monitor.stop()
Avalon.debug_info(_('Killing upscaled image cleaner'))
self.image_cleaner.stop()
def upscale(self,
input_directory,
output_directory,
scale_ratio,
processes,
push_strength=None,
push_grad_strength=None):
""" Anime4K wrapper
Arguments:
file_in {string} -- input file path
file_out {string} -- output file path
Keyword Arguments:
scale {int} -- scale ratio (default: {None})
push_strength {int} -- residual push strength (default: {None})
push_grad_strength {int} -- residual gradient push strength (default: {None})
Returns:
subprocess.Popen.returncode -- command line return value of execution
"""
# a list of all commands to be executed
commands = queue.Queue()
# get a list lof all image files in input_directory
extracted_frame_files = [
f for f in input_directory.iterdir() if
str(f).lower().endswith('.png') or str(f).lower().endswith('.jpg')
]
# upscale each image in input_directory
for image in extracted_frame_files:
execute = [
self.driver_settings['java_path'], '-jar',
self.driver_settings['path'],
str(image.absolute()),
str(output_directory / image.name),
str(scale_ratio)
]
# optional arguments
kwargs = ['push_strength', 'push_grad_strength']
# if optional argument specified, append value to execution list
for arg in kwargs:
if locals()[arg] is not None:
execute.extend([locals([arg])])
commands.put(execute)
# initialize two lists to hold running and finished processes
anime4k_running_processes = []
anime4k_finished_processes = []
# run all commands in queue
while not commands.empty():
# if any commands have completed
# remove the subprocess.Popen project and move it into finished processes
for process in anime4k_running_processes:
if process.poll() is not None:
Avalon.debug_info(
f'Subprocess {process.pid} exited with code {process.poll()}'
)
anime4k_finished_processes.append(process)
anime4k_running_processes.remove(process)
# when number running processes is less than what's specified
# create new processes and add to running process pool
while len(anime4k_running_processes) < processes:
next_in_queue = commands.get()
new_process = subprocess.Popen(next_in_queue)
anime4k_running_processes.append(new_process)
self.print_lock.acquire()
Avalon.debug_info(
f'[upscaler] Subprocess {new_process.pid} executing: {shlex.join(next_in_queue)}'
)
self.print_lock.release()
# return command execution return code
return anime4k_finished_processes
def _execute(self, execute):
# turn all list elements into string to avoid errors
execute = [str(e) for e in execute]
Avalon.debug_info(f'Executing: {shlex.join(execute)}')
return subprocess.Popen(execute)
def upscale(self, input_directory, output_directory, scale_ratio, upscaler_exceptions):
"""This is the core function for WAIFU2X class
Arguments:
input_directory {string} -- source directory path
output_directory {string} -- output directory path
ratio {int} -- output video ratio
"""
try:
# overwrite config file settings
self.waifu2x_settings['input_path'] = input_directory
self.waifu2x_settings['output_path'] = output_directory
# print thread start message
self.print_lock.acquire()
Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} started')
self.print_lock.release()
# waifu2x_ncnn_vulkan accepts arguments in a positional manner
# See: https://github.com/nihui/waifu2x_ncnn_vulkan#usage
# waifu2x_ncnn_vulkan.exe [input image] [output png] [noise=-1/0/1/2/3] [scale=1/2] [blocksize=400]
# noise = noise level, large value means strong denoise effect, -1=no effect
# scale = scale level, 1=no scale, 2=upscale 2x
# blocksize = tile size, use smaller value to reduce GPU memory usage, default is 400
# waifu2x_ncnn_vulkan does not accept an arbitrary scale ratio, max is 2
if scale_ratio == 1:
for raw_frame in os.listdir(input_directory):
command = [
os.path.join(input_directory, raw_frame),
os.path.join(output_directory, raw_frame),
str(self.waifu2x_settings['noise-level']),
'1',
str(self.waifu2x_settings['block-size'])
]
execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']]
execute.extend(command)
Avalon.debug_info(f'Executing: {execute}')
subprocess.run(execute, check=True, stderr=subprocess.DEVNULL)
else:
for raw_frame in os.listdir(input_directory):
command = [
os.path.join(input_directory, raw_frame),
os.path.join(output_directory, raw_frame),
str(self.waifu2x_settings['noise-level']),
'2',
str(self.waifu2x_settings['block-size'])
]
execute = [self.waifu2x_settings['waifu2x_ncnn_vulkan_path']]
execute.extend(command)
Avalon.debug_info(f'Executing: {execute}')
subprocess.run(execute, check=True, stderr=subprocess.DEVNULL)
# print thread exiting message
self.print_lock.acquire()
Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} exiting')
self.print_lock.release()
return 0
except Exception as e:
upscaler_exceptions.append(e)
# -------------------- Execution
if __name__ != "__main__":
Avalon.error("This file cannot be imported as a library")
raise ImportError("file cannot be imported")
# parse command line arguments
args = parse_arguments()
# print KPM icon
print_icon()
try:
# create KPM object
kpm = Kpm()
Avalon.debug_info("KPM initialized")
# privileged section
# check user privilege
if os.getuid() != 0:
Avalon.error("This program must be run as root")
sys.exit(1)
# if --install_kpm
if args.install_kpm:
# copy the current file to defined binary path
shutil.copy(pathlib.Path(__file__), KPM_PATH)
# change owner and permission of the file
os.chown(KPM_PATH, 0, 0)
KPM_PATH.chmod(0o755)
# print video2x logo
print_logo()
# parse command line arguments
video2x_args, driver_args = parse_arguments()
# display version and lawful informaition
if video2x_args.version:
print(LEGAL_INFO)
sys.exit(0)
# redirect output to both terminal and log file
if video2x_args.disable_logging is False:
LOGFILE = video2x_args.log
Avalon.debug_info(_('Redirecting console logs to {}').format(LOGFILE))
sys.stdout = BiLogger(sys.stdout, LOGFILE)
sys.stderr = BiLogger(sys.stderr, LOGFILE)
# read configurations from configuration file
config = read_config(video2x_args.config)
# load waifu2x configuration
driver_settings = config[video2x_args.driver]
driver_settings['path'] = os.path.expandvars(driver_settings['path'])
# read FFmpeg configuration
ffmpeg_settings = config['ffmpeg']
ffmpeg_settings['ffmpeg_path'] = os.path.expandvars(
ffmpeg_settings['ffmpeg_path'])
if args.force_upgrade:
Avalon.info('Force upgrading KPM from GitHub')
upgrade_kpm()
exit(0)
check_version()
# if -x, --xinstall specified
if args.xinstall:
packages = args.xinstall.split(',')
kobj.xinstall(packages)
# if no arguments are given
else:
kobj.upgrade_all()
Avalon.debug_info('Checking for unused packages')
# check if there are any unused packages
if kobj.autoremove_available():
if Avalon.ask('Remove useless packages?', True):
kobj.autoremove()
else:
Avalon.info('No unused packages found')
# apt autoclean
Avalon.info('Erasing old downloaded archive files')
kobj.autoclean()
except KeyboardInterrupt:
Avalon.warning('Aborting')
def _upscale_frames(self, w2):
""" Upscale video frames with waifu2x-caffe
This function upscales all the frames extracted
by ffmpeg using the waifu2x-caffe binary.
Arguments:
w2 {Waifu2x Object} -- initialized waifu2x object
"""
# progress bar thread exit signal
self.progress_bar_exit_signal = False
# create a container for exceptions in threads
# if this thread is not empty, then an exception has occured
self.upscaler_exceptions = []
# it's easier to do multi-threading with waifu2x_converter
# the number of threads can be passed directly to waifu2x_converter
if self.waifu2x_driver == 'waifu2x_converter':
progress_bar = threading.Thread(target=self._progress_bar,
args=([self.extracted_frames], ))
progress_bar.start()
w2.upscale(self.extracted_frames, self.upscaled_frames,
self.scale_ratio, self.threads,
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))
]:
renamed = re.sub('_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.png', '.png',
image)
shutil.move('{}\\{}'.format(self.upscaled_frames, image),
'{}\\{}'.format(self.upscaled_frames, renamed))
self.progress_bar_exit_signal = True
progress_bar.join()
return
# create a container for all upscaler threads
upscaler_threads = []
# list all images in the extracted frames
frames = [
os.path.join(self.extracted_frames, f)
for f in os.listdir(self.extracted_frames)
if os.path.isfile(os.path.join(self.extracted_frames, f))
]
# if we have less images than threads,
# create only the threads necessary
if len(frames) < self.threads:
self.threads = len(frames)
# create a folder for each thread and append folder
# name into a list
thread_pool = []
thread_folders = []
for thread_id in range(self.threads):
thread_folder = '{}\\{}'.format(self.extracted_frames,
str(thread_id))
thread_folders.append(thread_folder)
# delete old folders and create new folders
if os.path.isdir(thread_folder):
shutil.rmtree(thread_folder)
os.mkdir(thread_folder)
# append folder path into list
thread_pool.append((thread_folder, thread_id))
# evenly distribute images into each folder
# until there is none left in the folder
for image in frames:
# move image
shutil.move(image, thread_pool[0][0])
# rotate list
thread_pool = thread_pool[-1:] + thread_pool[:-1]
# create threads and start them
for thread_info in thread_pool:
# create thread
if self.scale_ratio:
thread = threading.Thread(target=w2.upscale,
args=(thread_info[0],
self.upscaled_frames,
self.scale_ratio, False, False,
self.upscaler_exceptions))
else:
thread = threading.Thread(
target=w2.upscale,
args=(thread_info[0], self.upscaled_frames, False,
self.scale_width, self.scale_height,
self.upscaler_exceptions))
thread.name = thread_info[1]
# add threads into the pool
upscaler_threads.append(thread)
# start progress bar in a different thread
progress_bar = threading.Thread(target=self._progress_bar,
args=(thread_folders, ))
progress_bar.start()
# create the clearer and start it
Avalon.debug_info('Starting upscaled image cleaner')
image_cleaner = ImageCleaner(self.extracted_frames,
self.upscaled_frames,
len(upscaler_threads))
image_cleaner.start()
# start all threads
for thread in upscaler_threads:
thread.start()
# wait for threads to finish
for thread in upscaler_threads:
thread.join()
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
if len(self.upscaler_exceptions) != 0:
raise (self.upscaler_exceptions[0])
def upgrade_all(self):
""" upgrade all packages
This method checks if there are packages
available for updating and update the packages
if updating them won't remove any packages from
the system. Often times when a bad source is added
to the system, APT tends to remove a number of packages
from the system when upgrading which is very risky.
"""
Avalon.info('Starting automatic upgrade')
Avalon.info('Updating APT cache')
with open('/etc/apt/sources.list', 'r') as aptlist:
for line in aptlist:
if 'ubuntu.com' in line and distro.linux_distribution(
)[0] != 'Ubuntu' and line.replace(' ', '')[0] != '#':
Avalon.warning('Ubuntu source detected in source.list!')
Avalon.warning(
'Continue upgrading might cause severe consequences!')
if Avalon.ask('Are you sure that you want to continue?',
False):
break
else:
Avalon.warning('Aborting system upgrade..')
exit(0)
self.update()
Avalon.info('APT cache updated')
if len(self.import_list) != 0:
if Avalon.ask('Detected unimported keys, import?', True):
if shutil.which('dirmngr') is None:
Avalon.warning('dirmngr Not installed')
Avalon.warning('It is required for importing keys')
# ask if user wants to install dirmngr
if Avalon.ask('Install Now?'):
self.install('dirnmgr')
# check dirmngr after package installation
if isinstance(shutil.which('dirmngr'), str):
Avalon.info('Installation successful')
self.import_keys(self.import_list)
Avalon.info('Keys imported')
Avalon.info(
'Updating APT cache after key importing')
self.update()
else:
Avalon.error('Installation Failed')
Avalon.error('Please check your settings')
Avalon.warning(
'dirmngr not available. Continuing without importing keys'
)
else:
Avalon.warning('dirmngr not available')
Avalon.warning('Continuing without importing keys')
else:
self.import_keys(self.import_list)
Avalon.info('Keys imported')
Avalon.info('Updating APT cache after key importing')
self.update()
# Second update after keys are imported
self.update()
# if there are no upgrades available
Avalon.debug_info('Checking package updates')
if self.no_upgrades():
Avalon.info('No upgrades available')
# if upgrades are available
else:
Avalon.debug_info('Checking if full upgrade is safe')
# if upgrade is safe, use -y flag on apt-get full-upgrade
# otherwise, let user confirm the upgrade
if self.full_upgrade_safe():
Avalon.info('Full upgrade is safe')
Avalon.info('Starting APT full upgrade')
self.full_upgrade()
else:
Avalon.warning('Full upgrade is NOT safe')
Avalon.warning('Requiring human confirmation')
self.manual_full_upgrade()
def wrapper(*args, **kwargs):
return function(*args, **kwargs)
Avalon.debug_info('{} row(s) affected'.format(args[0].cursor.rowcount))
def run(self):
""" Main controller for Video2X
This function controls the flow of video conversion
and handles all necessary functions.
"""
# external stop signal when called in a thread
self.running = True
# define process pool to contain processes
self.process_pool = []
# load driver modules
DriverWrapperMain = getattr(
importlib.import_module(f'wrappers.{self.driver}'), 'WrapperMain')
self.driver_object = DriverWrapperMain(self.driver_settings)
# load options from upscaler class into driver settings
self.driver_object.load_configurations(self)
# initialize FFmpeg object
self.ffmpeg_object = Ffmpeg(
self.ffmpeg_settings,
extracted_frame_format=self.extracted_frame_format)
# define processing queue
self.processing_queue = queue.Queue()
Avalon.info(_('Loading files into processing queue'))
Avalon.debug_info(_('Input path(s): {}').format(self.input))
# make output directory if the input is a list or a directory
if isinstance(self.input, list) or self.input.is_dir():
self.output.mkdir(parents=True, exist_ok=True)
input_files = []
# if input is single directory
# put it in a list for compability with the following code
if not isinstance(self.input, list):
input_paths = [self.input]
else:
input_paths = self.input
# flatten directories into file paths
for input_path in input_paths:
# if the input path is a single file
# add the file's path object to input_files
if input_path.is_file():
input_files.append(input_path)
# if the input path is a directory
# add all files under the directory into the input_files (non-recursive)
elif input_path.is_dir():
input_files.extend(
[f for f in input_path.iterdir() if f.is_file()])
output_paths = []
for input_path in input_files:
# get file type
# try python-magic if it's available
try:
input_file_mime_type = magic.from_file(str(
input_path.absolute()),
mime=True)
input_file_type = input_file_mime_type.split('/')[0]
input_file_subtype = input_file_mime_type.split('/')[1]
except Exception:
input_file_type = input_file_subtype = None
# in case python-magic fails to detect file type
# try guessing file mime type with mimetypes
if input_file_type not in ['image', 'video']:
input_file_mime_type = mimetypes.guess_type(input_path.name)[0]
input_file_type = input_file_mime_type.split('/')[0]
input_file_subtype = input_file_mime_type.split('/')[1]
# set default output file suffixes
# if image type is GIF, default output suffix is also .gif
if input_file_mime_type == 'image/gif':
output_path = self.output / self.output_file_name_format_string.format(
original_file_name=input_path.stem, extension='.gif')
elif input_file_type == 'image':
output_path = self.output / self.output_file_name_format_string.format(
original_file_name=input_path.stem,
extension=self.image_output_extension)
elif input_file_type == 'video':
output_path = self.output / self.output_file_name_format_string.format(
original_file_name=input_path.stem,
extension=self.video_output_extension)
# if file is none of: image, image/gif, video
# skip to the next task
else:
Avalon.error(
_('File {} ({}) neither an image nor a video').format(
input_path, input_file_mime_type))
Avalon.warning(_('Skipping this file'))
continue
# if there is only one input file
# do not modify output file suffix
if isinstance(self.input, pathlib.Path) and self.input.is_file():
output_path = self.output
output_path_id = 0
while str(output_path) in output_paths:
output_path = output_path.parent / pathlib.Path(
f'{output_path.stem}_{output_path_id}{output_path.suffix}')
output_path_id += 1
# record output path
output_paths.append(str(output_path))
# push file information into processing queue
self.processing_queue.put(
(input_path.absolute(), output_path.absolute(),
input_file_mime_type, input_file_type, input_file_subtype))
# check argument sanity before running
self._check_arguments()
# record file count for external calls
self.total_files = self.processing_queue.qsize()
Avalon.info(_('Loaded files into processing queue'))
# print all files in queue for debugging
for job in self.processing_queue.queue:
Avalon.debug_info(_('Input file: {}').format(job[0].absolute()))
try:
while not self.processing_queue.empty():
# get new job from queue
self.current_input_file, output_path, input_file_mime_type, input_file_type, input_file_subtype = self.processing_queue.get(
)
# start handling input
# if input file is a static image
if input_file_type == 'image' and input_file_subtype != 'gif':
Avalon.info(_('Starting to upscale image'))
self.process_pool.append(
self.driver_object.upscale(self.current_input_file,
output_path))
self._wait()
Avalon.info(_('Upscaling completed'))
# static images don't require GIF or video encoding
# go to the next task
self.processing_queue.task_done()
self.total_processed += 1
continue
# if input file is a image/gif file or a video
elif input_file_mime_type == 'image/gif' or input_file_type == 'video':
self.create_temp_directories()
# get video information JSON using FFprobe
Avalon.info(_('Reading video information'))
video_info = self.ffmpeg_object.probe_file_info(
self.current_input_file)
# analyze original video with FFprobe and retrieve framerate
# width, height = info['streams'][0]['width'], info['streams'][0]['height']
# find index of video stream
video_stream_index = None
for stream in video_info['streams']:
if stream['codec_type'] == 'video':
video_stream_index = stream['index']
break
# exit if no video stream found
if video_stream_index is None:
Avalon.error(_('Aborting: No video stream found'))
raise StreamNotFoundError('no video stream found')
# get average frame rate of video stream
framerate = float(
Fraction(video_info['streams'][video_stream_index]
['r_frame_rate']))
Avalon.info(_('Framerate: {}').format(framerate))
# self.ffmpeg_object.pixel_format = video_info['streams'][video_stream_index]['pix_fmt']
# extract frames from video
self.process_pool.append(
(self.ffmpeg_object.extract_frames(
self.current_input_file, self.extracted_frames)))
self._wait()
# if driver is waifu2x-caffe
# pass pixel format output depth information
if self.driver == 'waifu2x_caffe':
# get a dict of all pixel formats and corresponding bit depth
pixel_formats = self.ffmpeg_object.get_pixel_formats()
# try getting pixel format's corresponding bti depth
try:
self.driver_settings[
'output_depth'] = pixel_formats[
self.ffmpeg_object.pixel_format]
except KeyError:
Avalon.error(
_('Unsupported pixel format: {}').format(
self.ffmpeg_object.pixel_format))
raise UnsupportedPixelError(
f'unsupported pixel format {self.ffmpeg_object.pixel_format}'
)
# width/height will be coded width/height x upscale factor
# original_width = video_info['streams'][video_stream_index]['width']
# original_height = video_info['streams'][video_stream_index]['height']
# scale_width = int(self.scale_ratio * original_width)
# scale_height = int(self.scale_ratio * original_height)
# upscale images one by one using waifu2x
Avalon.info(_('Starting to upscale extracted frames'))
self._upscale_frames()
Avalon.info(_('Upscaling completed'))
# start handling output
# output can be either GIF or video
# if the desired output is gif file
if output_path.suffix.lower() == '.gif':
Avalon.info(
_('Converting extracted frames into GIF image'))
gifski_object = Gifski(self.gifski_settings)
self.process_pool.append(
gifski_object.make_gif(self.upscaled_frames,
output_path, framerate,
self.extracted_frame_format))
self._wait()
Avalon.info(_('Conversion completed'))
# if the desired output is video
else:
# frames to video
Avalon.info(_('Converting extracted frames into video'))
self.process_pool.append(
self.ffmpeg_object.assemble_video(
framerate, self.upscaled_frames))
# f'{scale_width}x{scale_height}'
self._wait()
Avalon.info(_('Conversion completed'))
try:
# migrate audio tracks and subtitles
Avalon.info(
_('Migrating audio, subtitles and other streams to upscaled video'
))
self.process_pool.append(
self.ffmpeg_object.migrate_streams(
self.current_input_file, output_path,
self.upscaled_frames))
self._wait()
# if failed to copy streams
# use file with only video stream
except subprocess.CalledProcessError:
traceback.print_exc()
Avalon.error(_('Failed to migrate streams'))
Avalon.warning(
_('Trying to output video without additional streams'
))
if input_file_mime_type == 'image/gif':
# copy will overwrite destination content if exists
shutil.copy(
self.upscaled_frames /
self.ffmpeg_object.intermediate_file_name,
output_path)
else:
# construct output file path
output_file_name = f'{output_path.stem}{self.ffmpeg_object.intermediate_file_name.suffix}'
output_video_path = output_path.parent / output_file_name
# if output file already exists
# create temporary directory in output folder
# temporary directories generated by tempfile are guaranteed to be unique
# and won't conflict with other files
if output_video_path.exists():
Avalon.error(_('Output video file exists'))
temporary_directory = pathlib.Path(
tempfile.mkdtemp(dir=output_path.parent))
output_video_path = temporary_directory / output_file_name
Avalon.info(
_('Created temporary directory to contain file'
))
# move file to new destination
Avalon.info(
_('Writing intermediate file to: {}').format(
output_video_path.absolute()))
shutil.move(
self.upscaled_frames /
self.ffmpeg_object.intermediate_file_name,
output_video_path)
# increment total number of files processed
self.cleanup_temp_directories()
self.processing_queue.task_done()
self.total_processed += 1
except (Exception, KeyboardInterrupt, SystemExit) as e:
with contextlib.suppress(ValueError, AttributeError):
self.cleanup_temp_directories()
self.running = False
raise e
# signal upscaling completion
self.running = False
def send_request(self):
"""send POST request to CloudFlare's server"""
global successes
global fails
global proxies
global thread_pool
succeeded = False
error_message = ""
try:
# generate random installation ID
install_id = "".join(
random.choices(string.ascii_uppercase + string.digits, k=22))
# construct POST request payload
post_data = {
"key":
f"{''.join(random.choices(string.ascii_uppercase + string.digits, k=43))}=",
"install_id": install_id,
"fcm_token":
f"{install_id}:APA91b{''.join(random.choices(string.ascii_uppercase + string.digits, k=134))}",
"referrer": self.warpid,
"warp_enabled": False,
"tos": datetime.datetime.utcnow().isoformat()[:-3] + "+00:00",
"type": "Android",
"locale": "en_US",
}
# send request
if proxies is None:
response = requests.post(
f"https://api.cloudflareclient.com/v0a{random.randint(100, 999)}/reg",
json=post_data,
headers=HEADERS,
timeout=self.timeout,
)
else:
proxy = proxies.popleft()
response = requests.post(
f"https://api.cloudflareclient.com/v0a{random.randint(100, 999)}/reg",
json=post_data,
headers=HEADERS,
timeout=self.timeout,
proxies={
"http": f"socks4://{proxy}",
"https": f"socks4://{proxy}",
},
)
if response.status_code == requests.codes.ok:
successes += 1
succeeded = True
else:
fails += 1
error_message = f" with code {response.status_code}"
if proxies:
proxies.append(proxy)
except IndexError:
self.running = False
return
# print error and carryon on upon exceptions
except Exception as e:
# _print(f"Thread {self.name} encountered an exception")
# _print(traceback.format_exc(), file=sys.stderr)
fails += 1
error_message = f" with Python exception {e}"
if proxies and not self.autoremove:
proxies.append(proxy)
finally:
information = [
f"Successes: {Avalon.FG.G}{successes}{Avalon.FG.DGR}",
f"Fails: {Avalon.FG.R}{fails}{Avalon.FG.DGR}",
# f"Live Threads: {threading.active_count()}",
f"Live Threads: {Avalon.FG.W}{len([t for t in thread_pool if t.is_alive()])}{Avalon.FG.DGR}",
]
# this information is not accurate
# information.append(f"Proxies in Pool: {len(proxies)}") if proxies else None
if succeeded:
Avalon.debug_info(
f"[{' | '.join(information)}] Thread {self.name} succeeded",
)
else:
Avalon.debug_info(
f"[{' | '.join(information)}] Thread {self.name} failed{error_message}",
)
def upscale(self, input_directory, output_directory, scale_ratio, jobs,
image_format, upscaler_exceptions):
""" Waifu2x Converter Driver Upscaler
This method executes the upscaling of extracted frames.
Arguments:
input_directory {string} -- source directory path
output_directory {string} -- output directory path
scale_ratio {int} -- frames' scale ratio
threads {int} -- number of threads
"""
try:
# overwrite config file settings
self.waifu2x_settings['input'] = input_directory
self.waifu2x_settings['output'] = output_directory
# temporary fix for https://github.com/DeadSix27/waifu2x-converter-cpp/issues/109
"""
self.waifu2x_settings['i'] = input_directory
self.waifu2x_settings['o'] = output_directory
self.waifu2x_settings['input'] = None
self.waifu2x_settings['output'] = None
"""
self.waifu2x_settings['scale-ratio'] = scale_ratio
self.waifu2x_settings['jobs'] = jobs
self.waifu2x_settings['output-format'] = image_format
# models_rgb must be specified manually for waifu2x-converter-cpp
# if it's not specified in the arguments, create automatically
if self.waifu2x_settings['model-dir'] is None:
self.waifu2x_settings['model-dir'] = pathlib.Path(
self.waifu2x_settings['waifu2x_converter_path']
) / 'models_rgb'
# print thread start message
self.print_lock.acquire()
Avalon.debug_info(
f'[upscaler] Thread {threading.current_thread().name} started')
self.print_lock.release()
# list to be executed
execute = []
for key in self.waifu2x_settings.keys():
value = self.waifu2x_settings[key]
# the key doesn't need to be passed in this case
if key == 'waifu2x_converter_path':
execute.append(
pathlib.Path(str(value)) / 'waifu2x-converter-cpp.exe')
# null or None means that leave this option out (keep default)
elif value is None or value is False:
continue
else:
if len(key) == 1:
execute.append(f'-{key}')
else:
execute.append(f'--{key}')
# true means key is an option
if value is True:
continue
execute.append(str(value))
Avalon.debug_info(f'Executing: {execute}')
return subprocess.run(execute, check=True).returncode
except Exception as e:
upscaler_exceptions.append(e)
def upscale(self, input_directory, output_directory, scale_ratio, jobs,
image_format):
""" Waifu2x Converter Driver Upscaler
This method executes the upscaling of extracted frames.
Arguments:
input_directory {string} -- source directory path
output_directory {string} -- output directory path
scale_ratio {int} -- frames' scale ratio
threads {int} -- number of threads
"""
# overwrite config file settings
self.driver_settings['input'] = input_directory
self.driver_settings['output'] = output_directory
self.driver_settings['scale-ratio'] = scale_ratio
self.driver_settings['jobs'] = jobs
self.driver_settings['output-format'] = image_format
# models_rgb must be specified manually for waifu2x-converter-cpp
# if it's not specified in the arguments, create automatically
if self.driver_settings['model-dir'] is None:
self.driver_settings['model-dir'] = pathlib.Path(
self.driver_settings['path']) / 'models_rgb'
# list to be executed
# initialize the list with waifu2x binary path as the first element
execute = [
str(
pathlib.Path(self.driver_settings['path']) /
'waifu2x-converter-cpp')
]
for key in self.driver_settings.keys():
value = self.driver_settings[key]
# the key doesn't need to be passed in this case
if key == 'path':
continue
# null or None means that leave this option out (keep default)
elif value is None or value is False:
continue
else:
if len(key) == 1:
execute.append(f'-{key}')
else:
execute.append(f'--{key}')
# true means key is an option
if value is True:
continue
execute.append(str(value))
# return the Popen object of the new process created
self.print_lock.acquire()
Avalon.debug_info(
f'[upscaler] Subprocess {os.getpid()} executing: {shlex.join(execute)}'
)
self.print_lock.release()
return subprocess.Popen(execute)
def _upscale_frames(self, input_directory: pathlib.Path,
output_directory: pathlib.Path):
"""Upscale video frames with waifu2x-caffe
This function upscales all the frames extracted
by ffmpeg using the waifu2x-caffe binary.
Args:
input_directory (pathlib.Path): directory containing frames to upscale
output_directory (pathlib.Path): directory which upscaled frames should be exported to
Raises:
UnrecognizedDriverError: raised when the given driver is not recognized
e: re-raised exception after an exception has been captured and finished processing in this scope
"""
# initialize waifu2x driver
if self.driver not in AVAILABLE_DRIVERS:
raise UnrecognizedDriverError(
_("Unrecognized driver: {}").format(self.driver))
# list all images in the extracted frames
frames = [(input_directory / f) for f in input_directory.iterdir()
if f.is_file]
# if we have less images than processes,
# create only the processes necessary
if len(frames) < self.processes:
self.processes = len(frames)
# create a directory for each process and append directory
# name into a list
process_directories = []
for process_id in range(self.processes):
process_directory = input_directory / str(process_id)
process_directories.append(process_directory)
# delete old directories and create new directories
if process_directory.is_dir():
shutil.rmtree(process_directory)
process_directory.mkdir(parents=True, exist_ok=True)
# waifu2x-converter-cpp will perform multi-threading within its own process
if self.driver in [
"waifu2x_converter_cpp",
"waifu2x_ncnn_vulkan",
"srmd_ncnn_vulkan",
"realsr_ncnn_vulkan",
"anime4kcpp",
]:
process_directories = [input_directory]
else:
# evenly distribute images into each directory
# until there is none left in the directory
for image in frames:
# move image
image.rename(process_directories[0] / image.name)
# rotate list
process_directories = (process_directories[-1:] +
process_directories[:-1])
# create driver processes and start them
for process_directory in process_directories:
self.process_pool.append(
self.driver_object.upscale(process_directory,
output_directory))
# start progress bar in a different thread
Avalon.debug_info(_("Starting progress monitor"))
self.progress_monitor = ProgressMonitor(self, process_directories)
self.progress_monitor.start()
# create the clearer and start it
Avalon.debug_info(_("Starting upscaled image cleaner"))
self.image_cleaner = ImageCleaner(input_directory, output_directory,
len(self.process_pool))
self.image_cleaner.start()
# wait for all process to exit
try:
self._wait()
except (Exception, KeyboardInterrupt, SystemExit) as e:
# cleanup
Avalon.debug_info(_("Killing progress monitor"))
self.progress_monitor.stop()
Avalon.debug_info(_("Killing upscaled image cleaner"))
self.image_cleaner.stop()
raise e
# if the driver is waifu2x-converter-cpp
# images need to be renamed to be recognizable for FFmpeg
if self.driver == "waifu2x_converter_cpp":
for image in [
f for f in output_directory.iterdir() if f.is_file()
]:
renamed = re.sub(
f"_\\[.*\\]\\[x(\\d+(\\.\\d+)?)\\]\\.{self.extracted_frame_format}",
f".{self.extracted_frame_format}",
str(image.name),
)
(output_directory / image).rename(output_directory / renamed)
# upscaling done, kill helper threads
Avalon.debug_info(_("Killing progress monitor"))
self.progress_monitor.stop()
Avalon.debug_info(_("Killing upscaled image cleaner"))
self.image_cleaner.stop()
def upscale(self, input_directory, output_directory, scale_ratio, upscaler_exceptions, push_strength=None, push_grad_strength=None):
""" Anime4K wrapper
Arguments:
file_in {string} -- input file path
file_out {string} -- output file path
Keyword Arguments:
scale {int} -- scale ratio (default: {None})
push_strength {int} -- residual push strength (default: {None})
push_grad_strength {int} -- residual gradient push strength (default: {None})
Returns:
subprocess.Popen.returncode -- command line return value of execution
"""
try:
# return value is the sum of all execution return codes
return_value = 0
# get a list lof all image files in input_directory
extracted_frame_files = [f for f in input_directory.iterdir() if str(f).lower().endswith('.png') or str(f).lower().endswith('.jpg')]
# Only print debug_info on first iteration.
logged = False
# upscale each image in input_directory
for image in extracted_frame_files:
execute = [
self.settings['java_path'] / self.settings['java_binary'],
'-jar',
self.settings['path'] / self.settings['binary'],
image.absolute(),
output_directory / image.name,
scale_ratio
]
# optional arguments
kwargs = [
'push_strength',
'push_grad_strength'
]
# if optional argument specified, append value to execution list
for arg in kwargs:
if locals()[arg] is not None:
execute.extend([locals([arg])])
# turn all list elements into string to avoid errors
execute = [str(e) for e in execute]
if not logged:
self.print_lock.acquire()
Avalon.debug_info(f'Executing: {execute}')
self.print_lock.release()
logged = True
return_value += subprocess.run(execute, check=True).returncode
# print thread exiting message
self.print_lock.acquire()
Avalon.debug_info(f'[upscaler] Thread {threading.current_thread().name} exiting')
self.print_lock.release()
# return command execution return code
return return_value
except Exception as e:
upscaler_exceptions.append(e)
def _upscale_frames(self):
""" Upscale video frames with waifu2x-caffe
This function upscales all the frames extracted
by ffmpeg using the waifu2x-caffe binary.
Arguments:
w2 {Waifu2x Object} -- initialized waifu2x object
"""
# progress bar thread exit signal
self.progress_bar_exit_signal = False
# create a container for exceptions in threads
# if this thread is not empty, then an exception has occured
self.upscaler_exceptions = []
# initialize waifu2x driver
drivers = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan']
if self.waifu2x_driver not in drivers:
raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
# it's easier to do multi-threading with waifu2x_converter
# the number of threads can be passed directly to waifu2x_converter
if self.waifu2x_driver == 'waifu2x_converter':
w2 = Waifu2xConverter(self.waifu2x_settings, self.model_dir)
progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],))
progress_bar.start()
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))]:
renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image)
shutil.move(os.path.join(self.upscaled_frames, image), os.path.join(self.upscaled_frames, renamed))
self.progress_bar_exit_signal = True
progress_bar.join()
return
else:
# create a container for all upscaler threads
upscaler_threads = []
# list all images in the extracted frames
frames = [os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f))]
# if we have less images than threads,
# create only the threads necessary
if len(frames) < self.threads:
self.threads = len(frames)
# create a directory for each thread and append directory
# name into a list
thread_pool = []
thread_directories = []
for thread_id in range(self.threads):
thread_directory = os.path.join(self.extracted_frames, str(thread_id))
thread_directories.append(thread_directory)
# delete old directories and create new directories
if os.path.isdir(thread_directory):
shutil.rmtree(thread_directory)
os.mkdir(thread_directory)
# append directory path into list
thread_pool.append((thread_directory, thread_id))
# evenly distribute images into each directory
# until there is none left in the directory
for image in frames:
# move image
shutil.move(image, thread_pool[0][0])
# rotate list
thread_pool = thread_pool[-1:] + thread_pool[:-1]
# create threads and start them
for thread_info in thread_pool:
# create a separate w2 instance for each thread
if self.waifu2x_driver == 'waifu2x_caffe':
w2 = Waifu2xCaffe(copy.deepcopy(self.waifu2x_settings), self.method, self.model_dir, self.bit_depth)
if self.scale_ratio:
thread = threading.Thread(target=w2.upscale,
args=(thread_info[0],
self.upscaled_frames,
self.scale_ratio,
False,
False,
self.image_format,
self.upscaler_exceptions))
else:
thread = threading.Thread(target=w2.upscale,
args=(thread_info[0],
self.upscaled_frames,
False,
self.scale_width,
self.scale_height,
self.image_format,
self.upscaler_exceptions))
# if the driver being used is waifu2x_ncnn_vulkan
elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan':
w2 = Waifu2xNcnnVulkan(copy.deepcopy(self.waifu2x_settings))
thread = threading.Thread(target=w2.upscale,
args=(thread_info[0],
self.upscaled_frames,
self.scale_ratio,
self.upscaler_exceptions))
# create thread
thread.name = thread_info[1]
# add threads into the pool
upscaler_threads.append(thread)
# start progress bar in a different thread
progress_bar = threading.Thread(target=self._progress_bar, args=(thread_directories,))
progress_bar.start()
# create the clearer and start it
Avalon.debug_info('Starting upscaled image cleaner')
image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
image_cleaner.start()
# start all threads
for thread in upscaler_threads:
thread.start()
# wait for threads to finish
for thread in upscaler_threads:
thread.join()
# upscaling done, kill the clearer
Avalon.debug_info('Killing upscaled image cleaner')
image_cleaner.stop()
self.progress_bar_exit_signal = True
if len(self.upscaler_exceptions) != 0:
raise(self.upscaler_exceptions[0])
# load cache directory
if config['video2x']['video2x_cache_directory'] is not None:
video2x_cache_directory = pathlib.Path(
config['video2x']['video2x_cache_directory'])
else:
video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x'
if video2x_cache_directory.exists() and not video2x_cache_directory.is_dir():
Avalon.error(_('Specified cache directory is a file/link'))
raise FileExistsError('Specified cache directory is a file/link')
# if cache directory doesn't exist
# ask the user if it should be created
elif not video2x_cache_directory.exists():
try:
Avalon.debug_info(
_('Creating cache directory {}').format(video2x_cache_directory))
video2x_cache_directory.mkdir(parents=True, exist_ok=True)
# there can be a number of exceptions here
# PermissionError, FileExistsError, etc.
# therefore, we put a catch-them-all here
except Exception as exception:
Avalon.error(_('Unable to create {}').format(video2x_cache_directory))
raise exception
# start execution
try:
# start timer
begin_time = time.time()
# if input specified is a single file
if video2x_args.input.is_file():
